Support multi-LoRA training with EP + FSDP2

cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py

@@ -0,0 +1,173 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""EP + FSDP2 + Multi-LoRA SFT cookbook for DeepSeek-V4.
+
+Run on 8 GPUs:
+    torchrun --nproc-per-node=8 cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py
+"""
+import os
+from pathlib import Path
+
+from peft import LoraConfig
+from transformers import AutoConfig
+
+import twinkle
+from twinkle import DeviceMesh, Platform, get_device_placement, get_logger
+from twinkle.dataloader import DataLoader
+from twinkle.dataset import Dataset, DatasetMeta
+from twinkle.model import MultiLoraTransformersModel
+from twinkle.preprocessor import SelfCognitionProcessor
+
+logger = get_logger()
+
+MODEL_ID = os.environ.get('DSV4_MODEL_ID', 'ms://deepseek-ai/DeepSeek-V4-Flash')
+DATASET_ID = os.environ.get('DATASET_ID', 'ms://swift/self-cognition')
+TEMPLATE_ID = os.environ.get('TEMPLATE_ID', 'DeepseekV4Template')
+BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '4'))
+GRAD_ACCUM_STEPS = int(os.environ.get('GRAD_ACCUM_STEPS', '4'))
+LOG_INTERVAL = GRAD_ACCUM_STEPS
+LR = float(os.environ.get('LR', '1e-4'))
+MAX_GRAD_NORM = float(os.environ.get('MAX_GRAD_NORM', '1.0'))
+LORA_R = int(os.environ.get('LORA_R', '8'))
+LORA_ALPHA = int(os.environ.get('LORA_ALPHA', '32'))
+MAX_LORAS = int(os.environ.get('MAX_LORAS', '2'))
+MAX_R = int(os.environ.get('MAX_R', str(max(32, LORA_R))))
+ENABLE_EP = os.environ.get('ENABLE_EP', '1') == '1'
+OUTPUT_DIR = os.environ.get('OUTPUT_DIR', './output_dsv4_multi_lora')
+RESUME_FROM_CHECKPOINT = os.environ.get('RESUME_FROM_CHECKPOINT') or None
+RESUME_ONLY_MODEL = os.environ.get('RESUME_ONLY_MODEL', '0') == '1'
+IGNORE_DATA_SKIP = os.environ.get('IGNORE_DATA_SKIP', '0') == '1'
+ADAPTER_NAMES = [name.strip() for name in os.environ.get('ADAPTER_NAMES', 'tenant_a,tenant_b').split(',') if name]
+
+device_mesh = DeviceMesh.from_sizes(
+    fsdp_size=8,
+    dp_size=1,
+    ep_size=8,
+    device_type=Platform.get_platform().device_prefix(),
+)
+twinkle.initialize(mode='local', global_device_mesh=device_mesh)
+
+
+def _build_lora_config(enable_ep: bool):
+    if enable_ep:
+        return LoraConfig(
+            r=LORA_R,
+            lora_alpha=LORA_ALPHA,
+            target_modules='all-linear',
+            exclude_modules=['o_a_proj'],
+            target_parameters=['mlp.experts.gate_up_proj', 'mlp.experts.down_proj'],
+        )
+    return LoraConfig(
+        r=LORA_R,
+        lora_alpha=LORA_ALPHA,
+        exclude_modules=['o_a_proj'],
+        target_modules='all-linear',
+    )
+
+
+def save_checkpoint(model: MultiLoraTransformersModel, adapter_name: str, dataloader: DataLoader):
+    return model.save(
+        name=f'checkpoint-final-{adapter_name}',
+        output_dir=OUTPUT_DIR,
+        adapter_name=adapter_name,
+        save_optimizer=True,
+        consumed_train_samples=dataloader.get_state()['consumed_train_samples'],
+    )
+
+
+def train():
+    config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True)
+    text_config = getattr(config, 'text_config', config)
+    if hasattr(text_config, 'use_cache'):
+        text_config.use_cache = False
+
+    dataset = Dataset(dataset_meta=DatasetMeta(DATASET_ID))
+    dataset.set_template(TEMPLATE_ID, model_id=MODEL_ID)
+    dataset.map(SelfCognitionProcessor('twinkle', 'ModelScope'))
+    dataset.encode(batched=True)
+    dataloader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, device_mesh=device_mesh)
+
+    ep_lora_cfg = _build_lora_config(enable_ep=ENABLE_EP) # LoraConfig for target params
+    lora_cfg = _build_lora_config(enable_ep=False)  # LoraConfig for PEFT adapter
+    model = MultiLoraTransformersModel(
+        model_id=MODEL_ID,
+        config=config,
+        device_mesh=device_mesh,
+        strategy='native_fsdp',
+        memory_efficient_init=True,
+        max_loras=MAX_LORAS,
+        max_r=MAX_R,
+        fsdp_config={
+            'expert_parallel': {
+                'enabled': ENABLE_EP,
+                'router_dtype': 'fp32',
+                'keep_router_logits': False,
+            }
+        },
+        lora_config=lora_cfg,
+    )
+
+    for adapter_name in ADAPTER_NAMES:
+        model.add_adapter_to_model(adapter_name, ep_lora_cfg, gradient_accumulation_steps=GRAD_ACCUM_STEPS)
+
+    if RESUME_FROM_CHECKPOINT:
+        checkpoint_path = Path(RESUME_FROM_CHECKPOINT).expanduser().resolve()
+        progress = None
+        for adapter_name in ADAPTER_NAMES:
+            progress = model.resume_from_checkpoint(
+                str(checkpoint_path),
+                resume_only_model=RESUME_ONLY_MODEL,
+                adapter_name=adapter_name,
+            )
+        if progress and not IGNORE_DATA_SKIP:
+            dataloader.resume_from_checkpoint(progress['consumed_train_samples'])
+
+    logger.info(get_device_placement())
+    for adapter_name in ADAPTER_NAMES:
+        logger.info(model.get_train_configs(adapter_name=adapter_name))
+    logger.info(
+        f'Total steps: {len(dataloader)}, batch_size={BATCH_SIZE}, grad_accum={GRAD_ACCUM_STEPS}, '
+        f'enable_ep={ENABLE_EP}, adapters={ADAPTER_NAMES}, output_dir={OUTPUT_DIR}')
+
+    # After LoRA init, before forward (LoRA active): perform EP + FSDP broadcast & sharding.
+    model._lazy_wrap_model()
+
+    # Must call set_optimizer() after EP + FSDP sharding, otherwise optimizer may
+    # capture stale parameter references and fail to update the actual LoRA weights.
+    for adapter_name in ADAPTER_NAMES:
+        model.set_optimizer('AdamW', lr=LR, foreach=False, adapter_name=adapter_name)
+        model.set_lr_scheduler(
+            scheduler_cls='CosineWarmupScheduler',
+            num_warmup_steps=5,
+            num_training_steps=len(dataloader),
+            adapter_name=adapter_name,
+        )
+
+    for adapter_name in ADAPTER_NAMES:
+        for batch_idx, batch in enumerate(dataloader):
+            if callable(batch):
+                batch = batch()
+            # adapter_name = ADAPTER_NAMES[batch_idx % len(ADAPTER_NAMES)]
+            model.forward_backward(
+                inputs=batch,
+                adapter_name=adapter_name,
+                gradient_accumulation_steps=GRAD_ACCUM_STEPS,
+            )
+            model.clip_grad_and_step(
+                max_grad_norm=MAX_GRAD_NORM,
+                adapter_name=adapter_name,
+                gradient_accumulation_steps=GRAD_ACCUM_STEPS,
+            )
+            cur_step = model.optimizer_group[adapter_name].cur_step
+            if cur_step > 0 and cur_step % LOG_INTERVAL == 0:
+                metric = model.calculate_metric(is_training=True, adapter_name=adapter_name)
+                if callable(metric):
+                    metric = metric()
+                logger.info(f'Adapter {adapter_name} is at step {cur_step} of {len(dataloader)}, metric: {metric}')
+
+    for adapter_name in ADAPTER_NAMES:
+        checkpoint = save_checkpoint(model, adapter_name, dataloader)
+        logger.info(f'Saved final adapter {adapter_name} to {checkpoint}')
+
+
+if __name__ == '__main__':
+    train()

src/twinkle/model/multi_lora.py

@@ -12,6 +12,7 @@
 from twinkle import torch_util
 from twinkle.data_format import InputFeature
 from twinkle.utils import get_logger
+from .multi_lora_target_parameters import TargetParameterLoraManager
 
 logger = get_logger()
 
@@ -36,6 +37,7 @@ def __init__(self, max_loras=5, max_r=32, max_length: int = 8192):
         self.module: PeftModel
         self._active_adapters = []
         self.max_length = max_length
+        self.target_parameter_manager = TargetParameterLoraManager(max_loras=max_loras, max_r=max_r)
 
     def _get_available_lora(self) -> Optional[LoraTenant]:
         for _lora in self.loras:
@@ -50,6 +52,10 @@ def _read_param_tensor(self, parameter):
     def _is_distributed_param(parameter):
         return hasattr(parameter, 'device_mesh') and hasattr(parameter, 'placements')
 
+    @staticmethod
+    def _is_target_parameter_lora_name(name: str) -> bool:
+        return '._twinkle_lora_' in name
+
     def _write_param_tensor(self, parameter, value):
         if value is None:
             return
@@ -142,15 +148,19 @@ def deactivate_adapter(self):
         else:
             self.module.disable_adapter_layers()
 
+    def patch_target_parameters(self, module, target_parameters):
+        self.target_parameter_manager.patch(module, target_parameters)
+
     @contextmanager
     def adapter(self, tenant_adapter_name: str, disable_lora: bool = False):
         self.activate_adapter(tenant_adapter_name)
-        if disable_lora:
-            # Temporarily disable all adapters while keeping optimizer_group active
-            with self._disable_lora_context(tenant_adapter_name):
+        with self.target_parameter_manager.adapter(tenant_adapter_name, disable_lora=disable_lora):
+            if disable_lora:
+                # Temporarily disable all adapters while keeping optimizer_group active
+                with self._disable_lora_context(tenant_adapter_name):
+                    yield self.find_lora_by_tenant(tenant_adapter_name).adapter_name
+            else:
                 yield self.find_lora_by_tenant(tenant_adapter_name).adapter_name
-        else:
-            yield self.find_lora_by_tenant(tenant_adapter_name).adapter_name
 
     @contextmanager
     def _disable_lora_context(self, tenant_adapter_name):
@@ -206,6 +216,12 @@ def acquire_lora(self, tenant_adapter_name: str, config: LoraConfig) -> str:
             raise RuntimeError(f'Too big rank for lora: {config.r}')
         _available_lora.tenant_config = config
         _available_lora.tenant_adapter_name = tenant_adapter_name
+        if getattr(config, 'target_parameters', None):
+            self.target_parameter_manager.acquire(
+                tenant_adapter_name=tenant_adapter_name,
+                slot_name=_available_lora.adapter_name,
+                config=config,
+            )
         logger.info(f'Lora count: {len(self.loras)}, available lora: {self._count_available_loras()}')
         return _available_lora.adapter_name
 
@@ -215,6 +231,7 @@ def release_lora(self, tenant_adapter_name: str) -> Optional[str]:
             _lora.tenant_config = None
             _lora.tenant_adapter_name = None
             self._load_initial_weights(_lora.adapter_name)
+            self.target_parameter_manager.release(tenant_adapter_name)
             logger.info(f'Lora count: {len(self.loras)}, available lora: {self._count_available_loras()}')
         except ValueError:
             return
@@ -462,20 +479,29 @@ def patch(self,
               target_modules='all-linear',
               *args,
               **kwargs):
+        module_device = getattr(module, 'device', None)
+        if module_device is None:
+            module_device = next(module.parameters())[1].device
+        low_cpu_mem_usage = module_device.type == 'meta'
+
         for i in range(self.max_loras):
-            config = LoraConfig(
-                r=self.max_r,
-                target_modules=target_modules,
-                lora_alpha=32,
-            )
+            config = kwargs.get('lora_config', None)
+            if config is None:
+                config = LoraConfig(
+                    r=self.max_r,
+                    target_modules=target_modules,
+                    lora_alpha=32,
+                    exclude_modules=['o_a_proj'],
+                )
             lora_tenant = LoraTenant(index=i, adapter_name=f'lora_{i}', config=config)
             self.loras.append(lora_tenant)
 
             def _patch_peft(_module):
                 if isinstance(_module, PeftModel):
-                    _module.add_adapter(lora_tenant.adapter_name, config)
+                    _module.add_adapter(lora_tenant.adapter_name, config, low_cpu_mem_usage=low_cpu_mem_usage)
                 else:
-                    _peft_model: PeftModel = get_peft_model(_module, config, lora_tenant.adapter_name)
+                    _peft_model: PeftModel = get_peft_model(
+                        _module, config, lora_tenant.adapter_name, low_cpu_mem_usage=low_cpu_mem_usage)
                     _module.active_adapters = _peft_model.active_adapters
                     _module = _peft_model
 
@@ -488,7 +514,7 @@ def _patch_megatron(_module):
                 # Expand target_modules (e.g., 'all-linear' -> actual module names)
                 _config = deepcopy(config)
                 if isinstance(_module, PeftModel):
-                    _module.add_adapter(lora_tenant.adapter_name, _config)
+                    _module.add_adapter(lora_tenant.adapter_name, _config, low_cpu_mem_usage=low_cpu_mem_usage)
                 else:
                     # TODO first wrap needs parse target_modules, need to fix later
                     if _config.target_modules:
@@ -499,7 +525,8 @@ def _patch_megatron(_module):
 
                         from .megatron import MegatronModel
                         _config.target_modules = MegatronModel.get_target_modules(_module, target_modules)
-                    _module = get_peft_model(_module, _config, lora_tenant.adapter_name)
+                    _module = get_peft_model(
+                        _module, _config, lora_tenant.adapter_name, low_cpu_mem_usage=low_cpu_mem_usage)
 
                 for name, submodule in _module.named_modules():
                     if isinstance(submodule, LoraLayer):
@@ -533,10 +560,12 @@ def save_initial_weights(self):
             lora_tenant = self.loras[i]
             pattern = re.compile(rf'\.lora_(?:A|embedding_A)\.{re.escape(lora_tenant.adapter_name)}\.')
 
-            def _store_weights(_module):
-                for name, parameter in _module.named_parameters():
-                    if pattern.search(name):
-                        lora_tenant.lora_A_weights[name] = self._read_param_tensor(parameter).clone().to('cpu')
+        def _store_weights(_module):
+            for name, parameter in _module.named_parameters():
+                if self._is_target_parameter_lora_name(name):
+                    continue
+                if pattern.search(name):
+                    lora_tenant.lora_A_weights[name] = self._read_param_tensor(parameter).clone().to('cpu')
 
             if isinstance(self.module, list):
                 for _module in self.module:
@@ -650,6 +679,8 @@ def set_state_dict(self, tenant_adapter_name, state_dict):
 
         def _load_weights(_module):
             for name, parameter in _module.named_parameters():
+                if self._is_target_parameter_lora_name(name):
+                    continue
                 if pattern.search(name) and self.match_target_modules(name, _lora.tenant_config.target_modules):
                     state_key = name.replace(f'.{_lora.adapter_name}.', '.')
                     target_tensor = self._read_param_tensor(parameter)
@@ -665,6 +696,7 @@ def _load_weights(_module):
                 _load_weights(_module)
         else:
             _load_weights(self.module)
+        self.target_parameter_manager.set_state_dict(tenant_adapter_name, state_dict)
 
     def get_state_dict(self, tenant_adapter_name):
         state_dict = {}
@@ -674,6 +706,8 @@ def get_state_dict(self, tenant_adapter_name):
         def _get_weights(_module):
             state_dict = {}
             for name, parameter in _module.named_parameters():
+                if self._is_target_parameter_lora_name(name):
+                    continue
                 if pattern.search(name) and self.match_target_modules(name, _lora.tenant_config.target_modules):
                     _param = self._slice_rank_tensor(name, self._read_param_tensor(parameter), _lora.tenant_config.r)
                     if _param is None:
@@ -687,6 +721,11 @@ def _get_weights(_module):
                 state_dict.update(_get_weights(_module))
         else:
             state_dict = _get_weights(self.module)
+        target_state_dict = self.target_parameter_manager.get_state_dict(tenant_adapter_name)
+        overlap = state_dict.keys() & target_state_dict.keys()
+        if overlap:
+            raise ValueError(f'Duplicate LoRA state keys: {sorted(overlap)[:5]}')
+        state_dict.update(target_state_dict)
         return state_dict
 
     def _load_initial_weights(self, origin_adapter_name):
@@ -696,6 +735,8 @@ def _load_initial_weights(self, origin_adapter_name):
 
         def _load_initial_weights(_module):
             for name, parameter in _module.named_parameters():
+                if self._is_target_parameter_lora_name(name):
+                    continue
                 if pattern_A.search(name):
                     local_param = self._read_param_tensor(parameter)
                     if local_param is not None:
@@ -794,3 +835,9 @@ def _get_parameters(_module):
         trainable_param_names = trainable_param_names[:5] + ['...'] + trainable_param_names[-5:]
         trainable_param_names = '\n'.join(trainable_param_names)
         return trainable_param_names
+
+    def get_target_parameter_trainable_parameters(self, tenant_adapter_name):
+        return {
+            name: parameter
+            for name, parameter in self.target_parameter_manager.named_slot_parameters(tenant_adapter_name)
+        }