wip

Author: tastelikefeet

cookbook/components/dataset.py

@@ -0,0 +1,4 @@
+from twinkle.dataset import Dataset
+
+
+dataset = Dataset('ms://swift/self-cognition')

cookbook/sft/lora_npu.py cookbook/npu/lora_npu.pycookbook/sft/lora_npu.py renamed to cookbook/npu/lora_npu.py

@@ -0,0 +1,99 @@
+import os
+
+from peft import LoraConfig
+from tqdm import tqdm
+
+import twinkle
+from twinkle import DeviceMesh, Platform
+from twinkle import get_device_placement, get_logger
+from twinkle.dataloader import DataLoader
+from twinkle.dataset import Dataset, DatasetMeta
+from twinkle.model import TransformersModel
+from twinkle.preprocessor import SelfCognitionProcessor
+
+if Platform.get_rank() == 0:
+    # rank0 recording
+    import swanlab
+    swanlab.login(api_key=os.environ['SWANLAB_API_KEY'], save=True)
+
+    run = swanlab.init(
+        project="megatron-swift",
+    )
+
+
+# Construct a device_mesh, fsdp=2, dp=2
+device_mesh = DeviceMesh.from_sizes(dp_size=2, fsdp_size=2)
+# use torchrun mode
+twinkle.initialize(mode='local', global_device_mesh=device_mesh)
+
+logger = get_logger()
+
+
+def eval(model):
+    # 100 Samples
+    dataset = Dataset(dataset_meta=DatasetMeta('ms://swift/self-cognition', data_slice=range(100)))
+    dataset.set_template('Template', model_id='ms://Qwen/Qwen2.5-7B-Instruct')
+    dataset.map(SelfCognitionProcessor('twinkle大模型', 'ModelScope社区'))
+    dataset.encode()
+    dataloader = DataLoader(dataset=dataset, batch_size=4)
+    for step, batch in tqdm(enumerate(dataloader)):
+        model.forward_only(inputs=batch)
+        model.calculate_loss()
+    metrics = model.calculate_metric(is_training=False)
+    return metrics
+
+
+def train():
+    # 1000 samples
+    dataset = Dataset(dataset_meta=DatasetMeta('ms://swift/self-cognition', data_slice=range(1000)))
+    # Set template to prepare encoding
+    dataset.set_template('Template', model_id='ms://Qwen/Qwen2.5-7B-Instruct')
+    # Preprocess the dataset to standard format
+    dataset.map(SelfCognitionProcessor('twinkle大模型', 'ModelScope社区'))
+    # Encode dataset
+    dataset.encode()
+    # Global batch size = 4, for GPUs, so 1 sample per GPU
+    dataloader = DataLoader(dataset=dataset, batch_size=4)
+    # Use a TransformersModel
+    model = TransformersModel(model_id='ms://Qwen/Qwen2.5-7B-Instruct')
+
+    lora_config = LoraConfig(
+        r=8,
+        lora_alpha=32,
+        target_modules='all-linear'
+    )
+
+    # Add a lora to model, with name `default`
+    model.add_adapter_to_model('default', lora_config, gradient_accumulation_steps=4)
+    # Add Optimizer for lora `default`
+    model.set_optimizer(optimizer_cls='AdamW', lr=1e-4)
+    # Add LRScheduler for lora `default`
+    model.set_lr_scheduler(scheduler_cls='CosineWarmupScheduler', num_warmup_steps=5, num_training_steps=len(dataloader))
+    logger.info(get_device_placement())
+    # Print the training config
+    logger.info(model.get_train_configs())
+    logger.info(f'Total steps: {len(dataloader)}')
+    loss_metric = 99.0
+    for step, batch in enumerate(dataloader):
+        # Do forward and backward
+        model.forward_backward(inputs=batch)
+        # Step
+        model.clip_grad_and_step()
+        if step % 20 == 0:
+            # Print metric
+            metric = model.calculate_metric(is_training=True)
+            if Platform.get_rank() == 0:
+                swanlab.log(metric)
+            logger.info(f'Current is step {step} of {len(dataloader)}, metric: {metric}')
+        if step > 0 and step % 40 == 0:
+           metrics = eval(model)
+           logger.info(f'Eval metric: {metrics}')
+           metrics['step'] = step
+           if loss_metric > float(metrics['loss']):
+               model.save(f'checkpoint-{step}')
+               loss_metric = float(metrics['loss'])
+    model.save(f'last-checkpoint', adapter_name='default')
+
+
+if __name__ == '__main__':
+    train()

cookbook/sampler_demo.py cookbook/sampler/sampler_demo.pycookbook/sampler_demo.py renamed to cookbook/sampler/sampler_demo.py

@@ -200,6 +200,7 @@ def __init__(self, # noqa
         self.sp_strategy = None
         self._model_wrapped = False
         self.optimizer_group: Dict[str, OptimizerGroup] = {_default_adapter_name: self._construct_default_optimizer_group()}
+        self.active_group = _default_adapter_name
 
     def _decide_strategy(self, strategy: Literal['accelerate', 'native_fsdp']):
         self._expert_parallel_config = self._fsdp_config.pop("expert_parallel", None)
@@ -243,11 +244,7 @@ def _get_default_group(self):
         """Get the only group has optimizer, else return the default one"""
         if len(self.optimizer_group) == 1:
             return next(iter(self.optimizer_group))
-        names = [name for name, og in self.optimizer_group.items() if og.optimizer is not None]
-        if names:
-            assert len(names) == 1, 'Only one group is supported.'
-            return names[0]
-        return _default_adapter_name
+        return self.active_group
 
     @staticmethod
     def _not_encoded(inputs):
@@ -905,6 +902,7 @@ def _patch_adapter(self, adapter_name: str, config_or_dir: Union[PeftConfig, str
         _gas_default = kwargs.get('gradient_accumulation_steps', 1)
         self.optimizer_group[adapter_name].gradient_accumulation_steps = _gas_default
         self._default_tokenizer = self.optimizer_group[adapter_name].template.processor
+        self.active_group = adapter_name
 
     @remote_function()
     def add_adapter_to_model(self, adapter_name: str, config_or_dir: Union[PeftConfig, str], **kwargs):