Qwen35 sp fix by ulysses

cookbook/transformers/sp_fsdp_dense.py

@@ -1,12 +1,15 @@
-import numpy as np
+import math
 from functools import partial
+
+import numpy as np
 from peft import LoraConfig
 
 import twinkle
-from twinkle import DeviceGroup, DeviceMesh, Platform, get_logger
+from twinkle import DeviceGroup, DeviceMesh, get_logger,Platform
 from twinkle.dataloader import DataLoader
 from twinkle.dataset import Dataset, DatasetMeta
 from twinkle.model import TransformersModel
+from twinkle.model.transformers.models import TwinkleQwen3_5ForCausalLM
 from twinkle.preprocessor import SelfCognitionProcessor
 
 logger = get_logger()
@@ -64,29 +67,41 @@ def train():
 
     model = TransformersModel(
         model_id=MODEL_ID,
+        model_cls=TwinkleQwen3_5ForCausalLM,
         device_mesh=device_mesh,
         strategy='native_fsdp',
+        attn_implementation='flash_attention_2'
     )
 
-    lora_config = LoraConfig(target_modules='all-linear')
-    model.add_adapter_to_model('default', lora_config, gradient_accumulation_steps=1)
+    lora_config = LoraConfig(target_modules='all-linear', lora_dropout=0.0)
+    model.add_adapter_to_model('default', lora_config)
+    grad_accumulation_steps = model.optimizer_group['default'].gradient_accumulation_steps
+    num_optimizer_steps = math.ceil(len(dataloader) / grad_accumulation_steps)
+    log_every_optimizer_steps = 20
     model.set_optimizer('AdamW', lr=1e-4, adapter_name='default')
     model.set_lr_scheduler(
         scheduler_cls='CosineWarmupScheduler',
         num_warmup_steps=5,
-        num_training_steps=len(dataloader),
+        num_training_steps=num_optimizer_steps,
         adapter_name='default',
     )
 
     logger.info(model.get_train_configs(adapter_name='default'))
-    logger.info(f'Total steps: {len(dataloader)}')
+    logger.info(
+        f'Total micro steps: {len(dataloader)}, optimizer steps: {num_optimizer_steps}, '
+        f'gradient_accumulation_steps: {grad_accumulation_steps}')
 
     for step, batch in enumerate(dataloader):
         model.forward_backward(inputs=batch, adapter_name='default')
         model.clip_grad_and_step(adapter_name='default')
-        if step % 20 == 0:
+        optimizer_step = step // grad_accumulation_steps
+        is_optimizer_boundary = (step + 1) % grad_accumulation_steps == 0
+        if is_optimizer_boundary and optimizer_step % log_every_optimizer_steps == 0:
             metric = model.calculate_metric(is_training=True, adapter_name='default')
-            logger.info(f'Current is step {step} of {len(dataloader)}, metric: {metric}')
+            optimizer_step = metric.get('iters', optimizer_step)
+            logger.info(
+                f'Current is optimizer step {optimizer_step} of {num_optimizer_steps} '
+                f'(micro step {step} of {len(dataloader)}), metric: {metric}')
     model.save('last-checkpoint', interval=1)
 
 

src/twinkle/dataloader/dataloader.py

@@ -45,14 +45,31 @@ def __init__(self,
         self.max_retries = kwargs.pop('max_retries', 20)
         self.min_batch_size = min_batch_size
         if device_mesh is not None:
-            assert batch_size >= device_mesh.data_world_size and batch_size % device_mesh.data_world_size == 0
-        self.batch_size = batch_size
+            required_world_size = self._required_data_world_size(device_mesh)
+            assert batch_size >= required_world_size and batch_size % required_world_size == 0
+        self.batch_size = self._resolve_runtime_batch_size(batch_size, device_mesh)
         self.dataloader_params = kwargs
-        self.dataloader_params['batch_size'] = batch_size
+        self.dataloader_params['batch_size'] = self.batch_size
         self.device_mesh = device_mesh
         self.processor: Optional[InputProcessor] = None
         self._set_work_init_fn()
 
+    @staticmethod
+    def _required_data_world_size(device_mesh: Optional[DeviceMesh]) -> int:
+        if device_mesh is None:
+            return 1
+        if (device_mesh.ulysses_size or 1) > 1:
+            return device_mesh.raw_data_world_size
+        return device_mesh.data_world_size
+
+    def _resolve_runtime_batch_size(self, batch_size: int, device_mesh: Optional[DeviceMesh]) -> int:
+        if device_mesh is None:
+            return batch_size
+        ulysses_size = device_mesh.ulysses_size or 1
+        if ulysses_size <= 1:
+            return batch_size
+        return batch_size // ulysses_size
+
     def _set_work_init_fn(self):
         num_workers = self.dataloader_params.get('num_workers', 2)
         self.dataloader_params['worker_init_fn'] = partial(

src/twinkle/model/transformers/__init__.py

@@ -1,3 +1,32 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from .multi_lora_transformers import MultiLoraTransformersModel
-from .transformers import TransformersModel
+from typing import TYPE_CHECKING
+
+from twinkle.utils.import_utils import _LazyModule
+
+if TYPE_CHECKING:
+    from .models import (TwinkleQwen3_5DecoderLayer, TwinkleQwen3_5ForCausalLM, TwinkleQwen3_5GatedDeltaNet,
+                         TwinkleQwen3_5PreTrainedModel, TwinkleQwen3_5TextModel)
+    from .multi_lora_transformers import MultiLoraTransformersModel
+    from .transformers import TransformersModel
+else:
+    _import_structure = {
+        'transformers': ['TransformersModel'],
+        'multi_lora_transformers': ['MultiLoraTransformersModel'],
+        'models': [
+            'TwinkleQwen3_5PreTrainedModel',
+            'TwinkleQwen3_5TextModel',
+            'TwinkleQwen3_5DecoderLayer',
+            'TwinkleQwen3_5GatedDeltaNet',
+            'TwinkleQwen3_5ForCausalLM',
+        ],
+    }
+
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,  # noqa
+        extra_objects={},
+    )

src/twinkle/model/transformers/models/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .qwen3_5 import (TwinkleQwen3_5DecoderLayer, TwinkleQwen3_5ForCausalLM, TwinkleQwen3_5GatedDeltaNet,
+                      TwinkleQwen3_5PreTrainedModel, TwinkleQwen3_5TextModel)
+
+__all__ = [
+    'TwinkleQwen3_5PreTrainedModel',
+    'TwinkleQwen3_5TextModel',
+    'TwinkleQwen3_5DecoderLayer',
+    'TwinkleQwen3_5GatedDeltaNet',
+    'TwinkleQwen3_5ForCausalLM',
+]

src/twinkle/model/transformers/models/qwen3_5/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .modeling_qwen3_5 import (TwinkleQwen3_5DecoderLayer, TwinkleQwen3_5ForCausalLM, TwinkleQwen3_5GatedDeltaNet,
+                               TwinkleQwen3_5PreTrainedModel, TwinkleQwen3_5TextModel)
+
+__all__ = [
+    'TwinkleQwen3_5PreTrainedModel',
+    'TwinkleQwen3_5TextModel',
+    'TwinkleQwen3_5DecoderLayer',
+    'TwinkleQwen3_5GatedDeltaNet',
+    'TwinkleQwen3_5ForCausalLM',
+]