diff --git a/cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py b/cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py new file mode 100644 index 00000000..0764dcda --- /dev/null +++ b/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() diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index d6032ade..23d0baa0 100644 --- a/src/twinkle/model/multi_lora.py +++ b/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) + } diff --git a/src/twinkle/model/multi_lora_target_parameters.py b/src/twinkle/model/multi_lora_target_parameters.py new file mode 100644 index 00000000..92da0dad --- /dev/null +++ b/src/twinkle/model/multi_lora_target_parameters.py @@ -0,0 +1,347 @@ +from __future__ import annotations + +import math +import torch +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from peft import LoraConfig +from torch import nn +from typing import Iterable, Iterator + + +class LoraParameterProxy(nn.Module): + + def __init__(self, LoraWrapper, slot_name): + super().__init__() + self.LoraWrapper = LoraWrapper + self.slot_name = slot_name + + def forward(self, weight: torch.Tensor) -> torch.Tensor: + delta_weight = self.LoraWrapper.get_delta_weight(self.slot_name) + mix_weight = weight + delta_weight + return mix_weight + + +@dataclass +class TargetParameterRecord: + module_name: str + module: nn.Module + parameter_name: str + + @property + def key(self) -> str: + if self.module_name: + return f'{self.module_name}.{self.parameter_name}' + return self.parameter_name + + +class TargetParameterLoraWrapper(nn.Module): + + def __init__(self, record: TargetParameterRecord, max_loras: int, max_r: int): + super().__init__() + self.record = record + # Unsharded original target parameter (pre-sharding snapshot) + parameter = getattr(self.record.module, self.record.parameter_name) + self.original_parameter_ndim = parameter.ndim + # Cache invariant attributes that won't change after initialization + self.did_swap_in_out_features = parameter.ndim == 3 and not getattr(self.record.module, 'is_transposed', False) + if parameter.ndim == 3: + self.in_features = parameter.shape[-1] if self.did_swap_in_out_features else parameter.shape[-2] + self.out_features = parameter.shape[-2] if self.did_swap_in_out_features else parameter.shape[-1] + else: + self.out_features, self.in_features = parameter.shape[0], parameter.shape[1] + + self.max_loras = max_loras + self.max_r = max_r + self.active_adapter: str | None = None + self.disable_adapters = False + self.lora_A = nn.ParameterDict() + self.lora_B = nn.ParameterDict() + self.scaling: dict[str, float] = {} + self.r: dict[str, int] = {} + self._initial_lora_A: dict[str, torch.Tensor] = {} + self.peft_key_prefix = '' + self._init_slots() + + @property + def base_parameter(self) -> nn.Parameter: + return getattr(self.record.module, self.record.parameter_name) + + @property + def num_experts(self) -> int: + # Check if the module has EP sharding info (for tensor experts) + if hasattr(self.record.module, '_ep_local_end') and hasattr(self.record.module, '_ep_local_start'): + return self.record.module._ep_local_end - self.record.module._ep_local_start # 在wrap_model的ep切分后会设置 + # Check if the module has EP sharding info (for ModuleList experts) + if hasattr(self.record.module, '_ep_experts_per_rank'): + return self.record.module._ep_experts_per_rank + # Fallback to base parameter shape + parameter = self.base_parameter + return parameter.shape[0] if parameter.ndim == 3 else 1 + + def _init_slots(self) -> None: + parameter = self.base_parameter + if parameter.ndim not in (2, 3): + raise ValueError( + f'target parameter {self.record.key} has {parameter.ndim} dimensions; only 2D and 3D are supported') + + # Note: reset_slot requires the tensor to be created on a physical device, not on a meta device. + device = parameter.device + if device.type == 'meta': + device = 'cpu' + for index in range(self.max_loras): + slot_name = f'lora_{index}' + self.lora_A[slot_name] = nn.Parameter( + torch.empty( + self.num_experts, + self.max_r, + self.in_features, + device=device, + dtype=parameter.dtype, + )) + self.lora_B[slot_name] = nn.Parameter( + torch.empty( + self.num_experts, + self.out_features, + self.max_r, + device=device, + dtype=parameter.dtype, + )) + self.r[slot_name] = self.max_r + self.scaling[slot_name] = 1.0 + self.reset_slot(slot_name) + + def reset_slot(self, slot_name: str) -> None: + if slot_name not in self._initial_lora_A: + nn.init.kaiming_uniform_(self.lora_A[slot_name], a=math.sqrt(5)) + self._initial_lora_A[slot_name] = self.lora_A[slot_name].detach().clone().cpu() + else: + initial = self._initial_lora_A[slot_name] + if hasattr(self.record.module, '_ep_local_start') and hasattr(self.record.module, '_ep_local_end'): + start = self.record.module._ep_local_start + end = self.record.module._ep_local_end + initial = initial[start:end] + initial = initial.to( + device=self.lora_A[slot_name].device, + dtype=self.lora_A[slot_name].dtype, + ) + self.lora_A[slot_name].data.copy_(initial) + nn.init.zeros_(self.lora_B[slot_name]) + + def configure_slot(self, slot_name: str, config: LoraConfig) -> None: + if slot_name not in self.lora_A: + raise ValueError(f'Unknown target-parameter LoRA slot: {slot_name}') + if config.r <= 0: + raise ValueError(f'`r` should be a positive integer value but the value passed is {config.r}') + if config.r > self.max_r: + raise ValueError(f'LoRA rank {config.r} exceeds max_r={self.max_r}') + if getattr(config, 'lora_dropout', 0): + raise ValueError('target_parameters LoRA does not support lora_dropout != 0') + if getattr(config, 'fan_in_fan_out', False): + raise ValueError('target_parameters LoRA does not support fan_in_fan_out=True') + if getattr(config, 'use_dora', False): + raise ValueError('target_parameters LoRA does not support use_dora=True') + if getattr(config, 'lora_bias', False): + raise ValueError('target_parameters LoRA does not support lora_bias=True') + + self.r[slot_name] = config.r + if getattr(config, 'use_rslora', False): + self.scaling[slot_name] = config.lora_alpha / math.sqrt(config.r) + else: + self.scaling[slot_name] = config.lora_alpha / config.r + + def get_delta_weight(self, slot_name: str) -> torch.Tensor: + if slot_name not in self.lora_A: + raise ValueError(f'Unknown target-parameter LoRA slot: {slot_name}') + + r = self.r[slot_name] + num_experts = self.num_experts + if hasattr(self.record.module, '_ep_local_start'): + # EP sharded: use full sharded weights + weight_A = self.lora_A[slot_name][:, :r, :] + weight_B = self.lora_B[slot_name][:, :, :r] + else: + # Not sharded: slice to actual rank + weight_A = self.lora_A[slot_name][:num_experts, :r, :] + weight_B = self.lora_B[slot_name][:num_experts, :, :r] + + # Don't call base_parameter during LoRA injection (infinite recursion risk) + if self.original_parameter_ndim == 2: + return (weight_B[0] @ weight_A[0]) * self.scaling[slot_name] + + if self.did_swap_in_out_features: + return torch.einsum('e o r, e r i -> e o i', weight_B, weight_A) * self.scaling[slot_name] + return torch.einsum('e o r, e r i -> e i o', weight_B, weight_A) * self.scaling[slot_name] + + @contextmanager + def activate(self, slot_name: str | None, disable_lora: bool = False): + if disable_lora or slot_name is None or slot_name not in self.lora_A: + yield + return + + module = self.record.module + param_name = self.record.parameter_name + already_parametrized = nn.utils.parametrize.is_parametrized(module, param_name) + if not already_parametrized: + # LoRA weights change after EP + FSDP sharding, so they must be computed dynamically and not be pre-fixed. + # delta_weight = self.get_delta_weight(slot_name) # lora_weight = B @ A * scaling + requires_grad_before = self.base_parameter.requires_grad + nn.utils.parametrize.register_parametrization( + self.record.module, + self.record.parameter_name, + LoraParameterProxy(self, slot_name), + ) + module.parametrizations[param_name].original.requires_grad_(requires_grad_before) + try: + with nn.utils.parametrize.cached(): + yield + finally: + if not already_parametrized: + nn.utils.parametrize.remove_parametrizations( + self.record.module, + self.record.parameter_name, + leave_parametrized=False, + ) + + def named_slot_parameters(self, slot_name: str) -> Iterator[tuple[str, nn.Parameter]]: + if slot_name not in self.lora_A: + return + yield f'{self.record.key}.lora_A.{slot_name}.weight', self.lora_A[slot_name] + yield f'{self.record.key}.lora_B.{slot_name}.weight', self.lora_B[slot_name] + + def parameters_for_slot(self, slot_name: str) -> list[nn.Parameter]: + return [parameter for _, parameter in self.named_slot_parameters(slot_name)] + + def get_state_dict(self, slot_name: str) -> dict[str, torch.Tensor]: + r = self.r[slot_name] + rank_width = r * self.num_experts + return { + f'{self.peft_key_prefix}.lora_A.weight': self.lora_A[slot_name][:rank_width, :].detach().clone(), + f'{self.peft_key_prefix}.lora_B.weight': self.lora_B[slot_name][:, :rank_width].detach().clone(), + } + + def set_state_dict(self, slot_name: str, state_dict: dict[str, torch.Tensor]) -> set[str]: + key_a = f'{self.peft_key_prefix}.lora_A.weight' + key_b = f'{self.peft_key_prefix}.lora_B.weight' + if key_a not in state_dict and key_b not in state_dict: + return set() + if key_a not in state_dict or key_b not in state_dict: + raise KeyError(f'Missing target-parameter LoRA pair for {self.peft_key_prefix}') + + with torch.no_grad(): + self.lora_A[slot_name].zero_() + self.lora_B[slot_name].zero_() + self.lora_A[slot_name][:state_dict[key_a].shape[0], :].copy_(state_dict[key_a].to( + device=self.lora_A[slot_name].device, + dtype=self.lora_A[slot_name].dtype, + )) + self.lora_B[slot_name][:, :state_dict[key_b].shape[1]].copy_(state_dict[key_b].to( + device=self.lora_B[slot_name].device, + dtype=self.lora_B[slot_name].dtype, + )) + return {key_a, key_b} + + +class TargetParameterLoraManager: + + def __init__(self, max_loras: int, max_r: int): + self.max_loras = max_loras + self.max_r = max_r + self.wrappers: list[TargetParameterLoraWrapper] = [] + self.tenant_to_slot: dict[str, str] = {} + self.tenant_configs: dict[str, LoraConfig] = {} + self._target_parameters: tuple[str, ...] | None = None + + def patch(self, model: nn.Module, target_parameters: Iterable[str]) -> None: + target_parameters = tuple(target_parameters) + if not target_parameters: + return + if self._target_parameters is not None: + if self._target_parameters != target_parameters: + raise ValueError( + f'target_parameters already patched as {self._target_parameters}, got {target_parameters}') + return + + records = [] + for module_name, module in model.named_modules(): + for param_name, parameter in module.named_parameters(recurse=False): + key = f'{module_name}.{param_name}' if module_name else param_name + if key in target_parameters or any(key.endswith(f'.{target}') for target in target_parameters): + if parameter.ndim not in (2, 3): + raise ValueError( + f'target parameter {key} has {parameter.ndim} dimensions; only 2D and 3D are supported') + records.append(TargetParameterRecord(module_name, module, param_name)) + + if not records: + raise ValueError(f'target_parameters={target_parameters} were set but no parameter was matched') + + for record in records: + wrapper = TargetParameterLoraWrapper(record, max_loras=self.max_loras, max_r=self.max_r) + record.module.add_module(f'_twinkle_lora_{record.parameter_name}', wrapper) + self.wrappers.append(wrapper) + self._assign_peft_key_prefixes() + self._target_parameters = target_parameters + + def _assign_peft_key_prefixes(self) -> None: + wrappers_by_module: dict[str, list[TargetParameterLoraWrapper]] = {} + for wrapper in self.wrappers: + wrappers_by_module.setdefault(wrapper.record.module_name, []).append(wrapper) + + for module_name, wrappers in wrappers_by_module.items(): + module_prefix = f'base_model.model.{module_name}' if module_name else 'base_model.model' + for index, wrapper in enumerate(wrappers): + nested_prefix = '.'.join(['base_layer'] * (len(wrappers) - index - 1)) + wrapper.peft_key_prefix = f'{module_prefix}.{nested_prefix}' if nested_prefix else module_prefix + + def acquire(self, tenant_adapter_name: str, slot_name: str, config: LoraConfig) -> None: + if tenant_adapter_name in self.tenant_to_slot: + raise ValueError(f'Lora {tenant_adapter_name} already exists') + if getattr(config, 'target_parameters', None) and not self.wrappers: + raise ValueError('target_parameters LoRA slots must be patched before acquire') + + self.tenant_to_slot[tenant_adapter_name] = slot_name + self.tenant_configs[tenant_adapter_name] = config + for wrapper in self.wrappers: + wrapper.configure_slot(slot_name, config) + + def release(self, tenant_adapter_name: str) -> None: + slot_name = self.tenant_to_slot.pop(tenant_adapter_name, None) + self.tenant_configs.pop(tenant_adapter_name, None) + if slot_name is None: + return + for wrapper in self.wrappers: + wrapper.reset_slot(slot_name) + + @contextmanager + def adapter(self, tenant_adapter_name: str, disable_lora: bool = False): + slot_name = self.tenant_to_slot.get(tenant_adapter_name) + with ExitStack() as stack: + for wrapper in self.wrappers: + stack.enter_context(wrapper.activate(slot_name, disable_lora=disable_lora)) + yield + + def parameters_for_tenant(self, tenant_adapter_name: str) -> list[nn.Parameter]: + slot_name = self.tenant_to_slot[tenant_adapter_name] + parameters = [] + for wrapper in self.wrappers: + parameters.extend(wrapper.parameters_for_slot(slot_name)) + return parameters + + def named_slot_parameters(self, tenant_adapter_name: str) -> Iterator[tuple[str, nn.Parameter]]: + slot_name = self.tenant_to_slot[tenant_adapter_name] + for wrapper in self.wrappers: + yield from wrapper.named_slot_parameters(slot_name) + + def get_state_dict(self, tenant_adapter_name: str) -> dict[str, torch.Tensor]: + slot_name = self.tenant_to_slot[tenant_adapter_name] + state_dict = {} + for wrapper in self.wrappers: + state_dict.update(wrapper.get_state_dict(slot_name)) + return state_dict + + def set_state_dict(self, tenant_adapter_name: str, state_dict: dict[str, torch.Tensor]) -> set[str]: + slot_name = self.tenant_to_slot[tenant_adapter_name] + consumed_keys = set() + for wrapper in self.wrappers: + consumed_keys.update(wrapper.set_state_dict(slot_name, state_dict)) + return consumed_keys diff --git a/src/twinkle/model/transformers/moe/expert_parallel.py b/src/twinkle/model/transformers/moe/expert_parallel.py index df3e9355..3562420d 100644 --- a/src/twinkle/model/transformers/moe/expert_parallel.py +++ b/src/twinkle/model/transformers/moe/expert_parallel.py @@ -142,6 +142,17 @@ def shard_experts( block._ep_tensor_experts = is_tensor_experts block._ep_ignore_shared_experts = cfg.ignore_shared_experts + # Duplicate EP metadata on block.experts as a defensive measure — + # this ensures the experts module has all needed context even if accessed separately. + block.experts._ep_num_experts = num_experts + block.experts._ep_experts_per_rank = experts_per_rank + block.experts._ep_local_start = local_start + block.experts._ep_local_end = local_end + block.experts._ep_rank = ep_rank + block.experts._ep_world_size = ep_world_size + block.experts._ep_tensor_experts = is_tensor_experts + block.experts._ep_ignore_shared_experts = cfg.ignore_shared_experts + return ExpertShardingSpec( block=block, experts_module=block.experts, @@ -441,6 +452,15 @@ def _shard_tensor_experts(experts: nn.Module, start: int, end: int) -> None: if hasattr(experts, 'num_experts'): experts.num_experts = end - start + from twinkle.model.multi_lora_target_parameters import TargetParameterLoraWrapper + for mod_name, target_param_wrapper in experts.named_children(): + if not isinstance(target_param_wrapper, TargetParameterLoraWrapper): + continue + for tenant_name, tenant_tensor in target_param_wrapper.lora_A.items(): + target_param_wrapper.lora_A[tenant_name] = nn.Parameter(tenant_tensor.data[start:end].clone()) + for tenant_name, tenant_tensor in target_param_wrapper.lora_B.items(): + target_param_wrapper.lora_B[tenant_name] = nn.Parameter(tenant_tensor.data[start:end].clone()) + def _run_local_experts( block: nn.Module, diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index 4c65807b..cd2fc924 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -32,6 +32,7 @@ def __init__( strategy: Literal['accelerate', 'native_fsdp'] = 'accelerate', ddp_config: Dict[str, Any] = None, fsdp_config: Dict[str, Any] = None, + lora_config: Dict[str, Any] = None, grad_scaler_config: Dict[str, Any] = None, memory_efficient_init: bool = False, max_loras: int = 5, @@ -49,6 +50,7 @@ def __init__( self._memory_efficient_init = memory_efficient_init self._decide_strategy(strategy) self.grad_scaler_config = grad_scaler_config + self.lora_config = lora_config or {} if model_id is not None: model_id = HubOperation.download_model(model_id) self.model_id = model_id @@ -65,7 +67,10 @@ def __init__( model_cls = getattr(transformers, model_cls) if model_id is None: self.model = model_cls.from_config(self.hf_config, **kwargs) + elif self._should_init_empty_pretrained_model_on_this_rank(): + self.model = self._init_empty_model_from_config(model_cls, **kwargs) else: + # Trigger transformers' FSDP-aware loading: meta-device init + rank-0-only weight load. with self.strategy.pretrained_load_context(): self.model = model_cls.from_pretrained(model_id, config=self.hf_config, **kwargs) self.tokenizer_id = kwargs.get('tokenizer_id', self.model_id) @@ -76,7 +81,7 @@ def __init__( self.optimizer_group: Dict[str, OptimizerGroup] = {} self.multi_adapter = MultiLora(max_loras=max_loras, max_r=max_r, max_length=max_length) self.model.gradient_checkpointing_enable() - self.model = self.multi_adapter.patch(self.model, target_modules=target_modules) + self.model = self.multi_adapter.patch(self.model, target_modules=target_modules, lora_config=self.lora_config) self.multi_adapter.save_initial_weights() # Active group for compatibility with single adapter self.active_group = None @@ -109,6 +114,22 @@ def unregister_mm_forward_hook(self, optimizer_group: OptimizerGroup): def _lazy_wrap_model(self): return super()._lazy_wrap_model() + def _maybe_apply_expert_parallel(self): + if self._memory_efficient_init: + raise NotImplementedError('Expert parallel is not supported with memory_efficient_init') + return super()._maybe_apply_expert_parallel() + + def _ensure_target_parameter_lora_installed(self, config: LoraConfig) -> None: + target_parameters = getattr(config, 'target_parameters', None) + if not target_parameters: + return + if self._model_wrapped: + raise RuntimeError('target_parameters LoRA must be installed before FSDP/DDP wrapping') + if getattr(self, '_enable_expert_parallel', False): + self.strategy.capture_pre_ep_state_if_needed(self.model, enable_ep=True) + # self._maybe_apply_expert_parallel() # 各rank广播之前不能对moe层进行分片, 没有实际权重时不能分片 + self.multi_adapter.patch_target_parameters(self.model, target_parameters) + @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) def forward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs): self._check_adapter_valid(kwargs.get('adapter_name')) @@ -201,6 +222,10 @@ def add_adapter_to_model(self, adapter_name: str, config_or_dir: Union[PeftConfi config_or_dir, str ), 'config_or_dir does not support str, because loading config from modelhub may causing unexpected behavior' assert isinstance(config_or_dir, LoraConfig), 'config_or_dir must be a LoraConfig instance' + config_or_dir = self.strategy.prepare_adapter_config( + config_or_dir, + enable_ep=getattr(self, '_enable_expert_parallel', False), + ) # Limit the max peft version in pyproject.toml, in case any newer version opens some untested module grad. config_or_dir.modules_to_save = None config_or_dir.bias = 'none' @@ -213,6 +238,7 @@ def add_adapter_to_model(self, adapter_name: str, config_or_dir: Union[PeftConfi _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._ensure_target_parameter_lora_installed(config_or_dir) self.multi_adapter.acquire_lora(tenant_adapter_name=adapter_name, config=config_or_dir) @remote_function() @@ -304,4 +330,9 @@ def _get_trainable_parameters_example(self, adapter_name, model): def _get_trainable_parameters(self, adapter_name): with self.multi_adapter.adapter(adapter_name) as real_adapter_name: - return super()._get_trainable_parameters(real_adapter_name) + params = super()._get_trainable_parameters(real_adapter_name) + # Note: experts have registered LoraWrapper as a submodule, so its internal LoRA parameters + # are already captured automatically. Duplicating parameter capture here will cause + # optimizer errors due to duplicate keys. + # params.update(self.multi_adapter.get_target_parameter_trainable_parameters(adapter_name)) + return params diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index 2722a900..5e515466 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -723,6 +723,10 @@ def _resolve_full_state_source_key(param_name: str, source_state: Mapping[str, A if _is_lora_state_key(param_name): raise KeyError(f"LoRA parameter '{param_name}' must be loaded from adapter source state.") + # Parametrization renames MoE layer parameters. full_sd stores the original + # weights, so we need to align the naming to match the expected keys. + if 'parametrizations' in param_name: + param_name = param_name.replace('.parametrizations', '').replace('.original', '') candidates = _source_key_candidates(param_name) for candidate in candidates: if candidate in source_state: diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index ee2ffe77..e0e28807 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -873,7 +873,9 @@ def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str, Optimizer], * def _get_trainable_parameters(self, adapter_name=_default_adapter_name): is_default = adapter_name == _default_adapter_name - pattern = re.compile(rf'\.lora_\w+\.{re.escape(adapter_name)}\.') + # Previously the pattern did not match nn.Parameter() weights, causing EP LoRA + # parameters to be missed by the optimizer. + pattern = re.compile(rf'\.lora_\w+\.{re.escape(adapter_name)}') params = {} model = self.strategy.unwrap_model(self.model) for name, param in model.named_parameters(): diff --git a/tests/model/test_multi_lora_target_parameters.py b/tests/model/test_multi_lora_target_parameters.py new file mode 100644 index 00000000..400e447b --- /dev/null +++ b/tests/model/test_multi_lora_target_parameters.py @@ -0,0 +1,256 @@ +import copy +import sys +import types + +import torch +from peft import LoraConfig, get_peft_model +from peft.utils import set_peft_model_state_dict +from torch import nn + + +class FakePackedExperts(nn.Module): + + def __init__(self, num_experts=2, hidden=4, intermediate=6, *, is_transposed=False): + super().__init__() + self.is_transposed = is_transposed + if is_transposed: + self.gate_up_proj = nn.Parameter(torch.randn(num_experts, intermediate * 2, hidden)) + self.down_proj = nn.Parameter(torch.randn(num_experts, hidden, intermediate)) + else: + self.gate_up_proj = nn.Parameter(torch.randn(num_experts, hidden, intermediate * 2)) + self.down_proj = nn.Parameter(torch.randn(num_experts, intermediate, hidden)) + + def forward(self, x, expert_idx=0): + gate_up = self.gate_up_proj[expert_idx] + down = self.down_proj[expert_idx] + if self.is_transposed: + hidden = torch.nn.functional.linear(x, gate_up) + gate, up = hidden.chunk(2, dim=-1) + return torch.nn.functional.linear(torch.nn.functional.silu(gate) * up, down) + + hidden = torch.nn.functional.linear(x, gate_up.T) + gate, up = hidden.chunk(2, dim=-1) + return torch.nn.functional.linear(torch.nn.functional.silu(gate) * up, down.T) + + +class FakeModel(nn.Module): + + def __init__(self, *, is_transposed=False): + super().__init__() + self.mlp = nn.Module() + self.mlp.experts = FakePackedExperts(is_transposed=is_transposed) + + def forward(self, x, expert_idx=0): + return self.mlp.experts(x, expert_idx=expert_idx) + + +def test_peft_target_parameter_key_shapes_for_3d_experts(): + model = FakeModel() + cfg = LoraConfig( + r=2, + lora_alpha=4, + target_modules=[], + target_parameters=["mlp.experts.gate_up_proj", "mlp.experts.down_proj"], + ) + + peft_model = get_peft_model(model, cfg, adapter_name="default") + state = peft_model.state_dict() + lora_shapes = {key: tuple(state[key].shape) for key in state if "lora_" in key} + + assert lora_shapes == { + "base_model.model.mlp.experts.base_layer.lora_A.default.weight": (4, 12), + "base_model.model.mlp.experts.base_layer.lora_B.default.weight": (4, 4), + "base_model.model.mlp.experts.lora_A.default.weight": (4, 4), + "base_model.model.mlp.experts.lora_B.default.weight": (6, 4), + } + + +def _make_target_cfg(r=2): + return LoraConfig( + r=r, + lora_alpha=r * 2, + target_modules=[], + target_parameters=["mlp.experts.gate_up_proj", "mlp.experts.down_proj"], + ) + + +def test_target_parameter_multi_lora_updates_only_active_adapter(): + from twinkle.model.multi_lora_target_parameters import TargetParameterLoraManager + + torch.manual_seed(0) + model = FakeModel() + manager = TargetParameterLoraManager(max_loras=2, max_r=4) + manager.patch(model, target_parameters=_make_target_cfg().target_parameters) + manager.acquire("adapter_a", "lora_0", _make_target_cfg(r=2)) + manager.acquire("adapter_b", "lora_1", _make_target_cfg(r=2)) + + params_before = { + name: param.detach().clone() + for name, param in manager.named_slot_parameters("adapter_b") + } + + opt = torch.optim.SGD(manager.parameters_for_tenant("adapter_a"), lr=0.1) + with manager.adapter("adapter_a"): + loss = model(torch.randn(3, 4), expert_idx=0).pow(2).mean() + loss.backward() + opt.step() + + for name, param in manager.named_slot_parameters("adapter_b"): + assert torch.equal(param.detach(), params_before[name]) + + +def test_multilora_releases_target_parameter_slot_to_initial_weights(): + from twinkle.model.multi_lora import LoraTenant, MultiLora + + torch.manual_seed(0) + model = FakeModel() + multi_lora = MultiLora(max_loras=2, max_r=4) + multi_lora.module = model + multi_lora.loras = [ + LoraTenant(index=0, adapter_name="lora_0", config=_make_target_cfg(r=4)), + LoraTenant(index=1, adapter_name="lora_1", config=_make_target_cfg(r=4)), + ] + multi_lora.patch_target_parameters(model, _make_target_cfg().target_parameters) + multi_lora.acquire_lora("adapter_a", _make_target_cfg(r=2)) + + initial_a = { + name: param.detach().clone() + for name, param in multi_lora.target_parameter_manager.named_slot_parameters("adapter_a") + if ".lora_A." in name + } + + with torch.no_grad(): + for _, param in multi_lora.target_parameter_manager.named_slot_parameters("adapter_a"): + param.add_(1.0) + + multi_lora.release_lora("adapter_a") + + for wrapper in multi_lora.target_parameter_manager.wrappers: + for name, param in wrapper.named_slot_parameters("lora_0"): + if ".lora_A." in name: + assert torch.equal(param.detach(), initial_a[name]) + else: + assert torch.count_nonzero(param.detach()) == 0 + + +def test_target_parameter_state_dict_loads_with_peft(): + from twinkle.model.multi_lora_target_parameters import TargetParameterLoraManager + + torch.manual_seed(0) + model = FakeModel() + base_state = copy.deepcopy(model.state_dict()) + cfg = _make_target_cfg(r=2) + manager = TargetParameterLoraManager(max_loras=2, max_r=4) + manager.patch(model, target_parameters=cfg.target_parameters) + manager.acquire("adapter_a", "lora_0", cfg) + + with torch.no_grad(): + for _, param in manager.named_slot_parameters("adapter_a"): + param.uniform_(-0.1, 0.1) + + inputs = torch.randn(3, 4) + with manager.adapter("adapter_a"): + expected = model(inputs, expert_idx=0) + state = manager.get_state_dict("adapter_a") + + peft_model = FakeModel() + peft_model.load_state_dict(base_state) + peft_model = get_peft_model(peft_model, cfg, adapter_name="default") + set_peft_model_state_dict(peft_model, state, adapter_name="default") + + actual = peft_model(inputs, expert_idx=0) + + assert torch.allclose(actual, expected, atol=1e-6) + + +def _make_multilora_for_target_parameters(model): + from twinkle.model.multi_lora import LoraTenant, MultiLora + + model.active_adapter = "lora_0" + model.set_adapter = lambda adapter_name: setattr(model, "active_adapter", adapter_name) + model.disable_adapter_layers = lambda: None + model.enable_adapter_layers = lambda: None + multi_lora = MultiLora(max_loras=2, max_r=4) + multi_lora.module = model + multi_lora.loras = [ + LoraTenant(index=0, adapter_name="lora_0", config=_make_target_cfg(r=4)), + LoraTenant(index=1, adapter_name="lora_1", config=_make_target_cfg(r=4)), + ] + multi_lora.patch_target_parameters(model, _make_target_cfg().target_parameters) + return multi_lora + + +def test_multilora_state_dict_round_trips_target_parameters(): + torch.manual_seed(0) + model = FakeModel() + base_state = copy.deepcopy(model.state_dict()) + multi_lora = _make_multilora_for_target_parameters(model) + multi_lora.acquire_lora("adapter_a", _make_target_cfg(r=2)) + + with torch.no_grad(): + for _, param in multi_lora.target_parameter_manager.named_slot_parameters("adapter_a"): + param.uniform_(-0.1, 0.1) + + inputs = torch.randn(3, 4) + with multi_lora.adapter("adapter_a"): + expected = model(inputs, expert_idx=0) + state = multi_lora.get_state_dict("adapter_a") + + restored_model = FakeModel() + restored_model.load_state_dict(base_state) + restored = _make_multilora_for_target_parameters(restored_model) + restored.acquire_lora("adapter_a", _make_target_cfg(r=2)) + restored.set_state_dict("adapter_a", state) + + with restored.adapter("adapter_a"): + actual = restored_model(inputs, expert_idx=0) + + assert torch.allclose(actual, expected, atol=1e-6) + + +def test_multilora_transformers_installs_target_parameters_once(): + from twinkle.model.multi_lora import LoraTenant, MultiLora + + if "zmq" not in sys.modules: + sys.modules["zmq"] = types.SimpleNamespace( + Context=object, + Socket=object, + RCVTIMEO=1, + SNDTIMEO=2, + LINGER=3, + ) + from twinkle.model.transformers.multi_lora_transformers import MultiLoraTransformersModel + + model = FakeModel() + instance = MultiLoraTransformersModel.__new__(MultiLoraTransformersModel) + nn.Module.__init__(instance) + instance.model = model + instance._model_wrapped = False + instance._enable_expert_parallel = False + instance.multi_adapter = MultiLora(max_loras=2, max_r=4) + instance.multi_adapter.module = model + instance.multi_adapter.loras = [ + LoraTenant(index=0, adapter_name="lora_0", config=_make_target_cfg(r=4)), + LoraTenant(index=1, adapter_name="lora_1", config=_make_target_cfg(r=4)), + ] + + cfg = _make_target_cfg(r=2) + instance._ensure_target_parameter_lora_installed(cfg) + instance.multi_adapter.acquire_lora("adapter_a", cfg) + instance._ensure_target_parameter_lora_installed(cfg) + instance.multi_adapter.acquire_lora("adapter_b", cfg) + + assert len(instance.multi_adapter.target_parameter_manager.wrappers) == 2 + + different_cfg = LoraConfig( + r=2, + lora_alpha=4, + target_modules=[], + target_parameters=["mlp.experts.gate_up_proj"], + ) + try: + instance._ensure_target_parameter_lora_installed(different_cfg) + except ValueError: + pass + else: + raise AssertionError("different target_parameters should be rejected") diff --git a/tests/moe/test_ep_multi_lora_target_parameters.py b/tests/moe/test_ep_multi_lora_target_parameters.py new file mode 100644 index 00000000..91fdb718 --- /dev/null +++ b/tests/moe/test_ep_multi_lora_target_parameters.py @@ -0,0 +1,31 @@ +import pytest +import sys +import torch +import types + + +def _ensure_dummy_zmq(): + if "zmq" in sys.modules: + return + sys.modules["zmq"] = types.SimpleNamespace( + Context=object, + Socket=object, + RCVTIMEO=1, + SNDTIMEO=2, + LINGER=3, + ) + + +def test_ep_target_parameter_lora_gather_dim_matches_peft_flattening(): + _ensure_dummy_zmq() + from twinkle.model.transformers.strategy.native_fsdp import _ep_expert_state_dict_gather_dim + + assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.lora_A.weight") == 0 + assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.base_layer.lora_A.weight") == 0 + assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.lora_B.weight") == 1 + assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.base_layer.lora_B.weight") == 1 + + +@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() < 4, reason="Need 4 GPUs") +def test_ep_fsdp_multi_lora_target_parameter_checkpoint_smoke(): + pytest.skip("Run this smoke in the DSV4 EP/FSDP integration environment with a local model fixture.")