# Copyright (c) ModelScope Contributors. All rights reserved. import re import torch import torch.distributed as dist from megatron.core import mpu from megatron.core.extensions.transformer_engine import TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear from megatron.core.inference.communication_utils import recv_from_prev_pipeline_rank_, send_to_next_pipeline_rank from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing from megatron.core.transformer.moe.router import TopKRouter from torch import nn from transformers.utils import is_torch_npu_available from swift.tuners import LoraConfig, Swift from swift.utils import (activate_parameters, deep_getattr, find_layers, freeze_parameters, get_logger, get_model_parameter_info) from swift.utils import get_packed_seq_params as _get_packed_seq_params logger = get_logger() def find_all_linears(model, extra_layers=None): def _cond(name, module): if (extra_layers and isinstance(module, tuple(extra_layers))) or name != 'output_layer' and isinstance( module, (TELinear, TELayerNormColumnParallelLinear, TEGroupedLinear, nn.Linear)): return True return False return find_layers(model, _cond) def find_router(model): return find_layers(model, lambda name, module: isinstance(module, TopKRouter)) def find_embedding(model): return find_layers(model, lambda name, module: isinstance(module, LanguageModelEmbedding)) def get_multimodal_target_regex( args, model, *, freeze_llm: bool = False, freeze_vit: bool = True, freeze_aligner: bool = True, include_embedding: bool = False, include_router: bool = False, ) -> str: megatron_model_meta = args.megatron_model_meta modules = [] visual_cls = megatron_model_meta.visual_cls vision_tower = [f'visual.{vit}' for vit in visual_cls._vision_tower] aligner = [f'visual.{aligner}' for aligner in visual_cls._aligner] if not freeze_llm: modules.append('language_model') if not freeze_vit: modules += vision_tower if not freeze_aligner: modules += aligner assert len(modules) > 0, f'modules: {modules}' extra_layers = [] if include_embedding: extra_layers.append(LanguageModelEmbedding) if include_router: extra_layers.append(TopKRouter) res = [] for module in modules: rejected_modules = [] if not freeze_vit: for _aligner in aligner: if _aligner.startswith(f'{module}.'): rejected_modules.append(_aligner) sub_module = deep_getattr(model, module) if sub_module is None: continue target_modules = find_all_linears(sub_module, extra_layers) if not target_modules: continue target_modules = [tm for tm in target_modules if tm] target_pattern = rf'.*\.({"|".join(target_modules)})' if target_modules else '' rejected_pattern = rf'(?!({"|".join(rejected_modules)}))' if rejected_modules else '' res.append(rf'{rejected_pattern}{re.escape(module)}(?=\.){target_pattern}') return rf'^({"|".join(res)})$' def get_target_modules(args, model): if isinstance(args.target_modules, str): return args.target_modules target_modules = args.target_modules.copy() if 'all-linear' in target_modules: if args.is_multimodal and not args.language_model_only: if args.tuner_type == 'lora_llm': kwargs = { 'freeze_llm': False, 'freeze_vit': True, 'freeze_aligner': True, } else: # lora kwargs = { 'freeze_llm': args.freeze_llm, 'freeze_vit': args.freeze_vit, 'freeze_aligner': args.freeze_aligner, } return get_multimodal_target_regex( args, model, include_embedding='all-embedding' in target_modules, include_router='all-router' in target_modules, **kwargs, ) else: target_modules.remove('all-linear') target_modules += find_all_linears(model) if 'all-embedding' in target_modules: target_modules.remove('all-embedding') target_modules += find_embedding(model) if 'all-router' in target_modules: target_modules.remove('all-router') target_modules += find_router(model) return target_modules def get_modules_to_save(args, model): if args.task_type == 'seq_cls': args.modules_to_save.append('output_layer') modules_to_save = args.modules_to_save.copy() if 'all-embedding' in args.modules_to_save: modules_to_save.remove('all-embedding') modules_to_save += find_embedding(model) return modules_to_save def prepare_adapter(args, model): target_modules = get_target_modules(args, model) modules_to_save = get_modules_to_save(args, model) lora_kwargs = { 'r': args.lora_rank, 'target_modules': target_modules, 'lora_alpha': args.lora_alpha, 'lora_dropout': args.lora_dropout, 'bias': args.lora_bias, 'modules_to_save': modules_to_save, 'use_rslora': args.use_rslora, } lora_config = LoraConfig(task_type='CAUSAL_LM', lora_dtype=args.lora_dtype, **lora_kwargs) logger.info(f'lora_config: {lora_config}') model = Swift.prepare_model(model, lora_config) if args.mcore_ref_adapter or args.ref_adapters: model.add_adapter('ref_adapter', lora_config) model.base_model._cast_adapter_dtype(adapter_name='ref_adapter', autocast_adapter_dtype=True) for n, p in model.named_parameters(): if '.ref_adapter.' in n: p.requires_grad = False return model def _prepare_full_vit(args, model): megatron_model_meta = args.megatron_model_meta visual_cls = megatron_model_meta.visual_cls vision_tower = [f'visual.{vit}' for vit in visual_cls._vision_tower] aligner = [f'visual.{aligner}' for aligner in visual_cls._aligner] for module_prefix in vision_tower + aligner: module = deep_getattr(model, module_prefix) if module is not None: module.requires_grad_(True) def prepare_mcore_model(args, model): if args.tuner_type == 'full': freeze_parameters(model, args.freeze_parameters_ratio, args.freeze_parameters, args.freeze_parameters_regex) if args.trainable_parameters or args.trainable_parameters_regex: activate_parameters(model, args.trainable_parameters, args.trainable_parameters_regex) elif args.tuner_type in {'lora', 'lora_llm'}: model = prepare_adapter(args, model) if args.tuner_type == 'lora_llm': _prepare_full_vit(args, model) logger.info(f'model: {model}') logger.info_if( f'[rank{dist.get_rank()}] model_parameter_info: {get_model_parameter_info(model)}', cond=mpu.get_data_parallel_rank() == 0) return model def forward_step_helper(model, inputs, dtype=None): config = model.config dtype = dtype or config.params_dtype if not mpu.is_pipeline_first_stage(): recv_shape_buffer = torch.empty((3, ), device=torch.cuda.current_device(), dtype=torch.int64) recv_from_prev_pipeline_rank_(recv_shape_buffer) recv_buffer = torch.empty(recv_shape_buffer.tolist(), device=torch.cuda.current_device(), dtype=dtype) recv_from_prev_pipeline_rank_(recv_buffer) model.set_input_tensor(recv_buffer) output_tensor = model(**inputs) if not mpu.is_pipeline_last_stage(): recv_shape_buffer = torch.tensor(output_tensor.shape, device=torch.cuda.current_device(), dtype=torch.int64) send_to_next_pipeline_rank(recv_shape_buffer) send_to_next_pipeline_rank(output_tensor) output_tensor = None return output_tensor def get_padding_to(args): padding_to = None if args.tensor_model_parallel_size > 1 and args.sequence_parallel: padding_to = args.tensor_model_parallel_size if args.context_parallel_size > 1: padding_to = (padding_to or 1) * args.context_parallel_size origin_padding_to = padding_to fp8_format = getattr(args, 'fp8_format', None) or getattr(args, 'fp8', None) fp4_format = getattr(args, 'fp4_format', None) or getattr(args, 'fp4', None) if args.fp8_recipe == 'blockwise': padding_to = (padding_to or 1) * 128 elif args.fp8_recipe == 'mxfp8': # MXFP8 uses a block size of 32. Under sequence parallel, the sequence is # split across TP ranks, so each per-rank shard (seq_len / TP) must itself # be divisible by 32. Pad the total length to TP * 32 to guarantee this. padding_to = (padding_to or 1) * 32 elif fp8_format is not None or fp4_format is not None: padding_to = (padding_to or 1) * 16 if args.attention_backend == 'fused': padding_to = max(padding_to or 1, ((origin_padding_to) or 1) * 64) return padding_to def get_packed_seq_params(args, position_ids: torch.Tensor) -> PackedSeqParams: params = _get_packed_seq_params(position_ids) packed = PackedSeqParams( cu_seqlens_q=params['cu_seq_lens_q'], cu_seqlens_kv=params['cu_seq_lens_k'], max_seqlen_q=params['max_length_q'], max_seqlen_kv=params['max_length_k'], qkv_format='thd', ) if hasattr(packed, 'cp_partition_mode'): packed.cp_partition_mode = args.cp_partition_mode if is_torch_npu_available(): packed.cu_seqlens_q_padded = params['cu_seq_lens_q'] packed.cu_seqlens_kv_padded = params['cu_seq_lens_k'] return packed def reconstruct_tensor_cp(tensor, packed_seq_params, dim=1) -> torch.Tensor: """In CP mode, all-gather and undo the load-balanced (zigzag) chunking produced by ``split_cp_inputs``, restoring the full sequence in original token order along ``dim``. Args: tensor: CP-sharded local tensor whose sequence dim is at ``dim``. packed_seq_params: ``PackedSeqParams`` for THD inputs, or ``None`` for regular ``[B, S, ...]`` inputs. dim: Sequence dimension index of ``tensor`` (default: 1). Returns: torch.Tensor: Full-sequence tensor with the same shape as ``tensor`` except the size at ``dim`` is multiplied by ``cp_size``. """ cp_size = mpu.get_context_parallel_world_size() if cp_size <= 1: return tensor cp_rank = mpu.get_context_parallel_rank() cp_group = mpu.get_context_parallel_group() # All-gather across CP ranks (preserve local autograd graph for `tensor`). output_list = [torch.empty_like(tensor) for _ in range(cp_size)] torch.distributed.all_gather(output_list, tensor.contiguous(), group=cp_group) output_list[cp_rank] = tensor gathered = torch.cat(output_list, dim=dim) # `_undo_attention_load_balancing` assumes sequence dim is 0; transpose if needed. if dim != 0: gathered = gathered.transpose(0, dim).contiguous() out = _undo_attention_load_balancing(gathered, cp_size, packed_seq_params) if dim != 0: out = out.transpose(0, dim).contiguous() return out