code/xtuner/model/llava.py

# Copyright (c) OpenMMLab. All rights reserved.
import math
import os.path as osp
import warnings
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.distributed as dist  # === MOD ===
from accelerate import init_empty_weights
from mmengine import print_log
from mmengine.config import Config, ConfigDict
from mmengine.model import BaseModel
from peft import get_peft_model, prepare_model_for_kbit_training
from transformers import (AddedToken, AutoConfig, CLIPImageProcessor,
                          CLIPVisionModel, LlamaForCausalLM,
                          LlamaTokenizerFast, LlavaConfig,
                          LlavaForConditionalGeneration, LlavaProcessor)
from transformers.integrations import is_deepspeed_zero3_enabled
import os
from safetensors.torch import load_file, save_file
from xtuner.registry import BUILDER
from xtuner.utils import DEFAULT_IMAGE_TOKEN
from .modules import ProjectorConfig, ProjectorModel, dispatch_modules
from .modules.dispatch import SUPPORT_FLASH1, SUPPORT_FLASH2
from .utils import (LoadWoInit, find_all_linear_names,
                    get_peft_model_state_dict, guess_load_checkpoint,
                    make_inputs_require_grad,
                    prepare_inputs_labels_for_multimodal, traverse_dict)

from .torchscale.model.LongNetWithMerging import make_swin_longnet_from_name
# from xtuner.model.torchscale.model.LongNet import make_longnet_from_name
from .torchscale.model.LongNet import make_longnet_from_name
import torch.nn.functional as F

# ===== 在类前或类内其它位置都可以：新增一个探测函数 =====
def _detect_qwen_major_version(llm) -> int:
    """
    返回 3 表示 Qwen3，2 表示 Qwen2，0 表示未知/其它。
    优先用 config.model_type，其次回退到类名字符串。
    """
    base = llm.model if hasattr(llm, "model") else llm
    cfg = getattr(base, "config", None)
    mt = (getattr(cfg, "model_type", None) or "").lower()
    if mt == "qwen3":
        return 3
    if mt == "qwen2":
        return 2

    # 回退：根据类名判别
    cname = base.__class__.__name__.lower()
    if "qwen3" in cname:
        return 3
    if "qwen2" in cname:
        return 2
    return 0

def convert_state_dict_to_hf(state_dict, mapping):
    new_state_dict = {}
    for key, value in state_dict.items():
        if key.endswith('.inv_freq'):
            continue
        for key_to_modify, new_key in mapping.items():
            if key_to_modify in key:
                key = key.replace(key_to_modify, new_key)
        new_state_dict[key] = value
    return new_state_dict

class AdaptiveAvgPool1dLayer(nn.Module):
    def __init__(self, output_size):
        super(AdaptiveAvgPool1dLayer, self).__init__()
        self.output_size = output_size

    def forward(self, x):
        return F.adaptive_avg_pool1d(x, self.output_size)
    
class LLaVAModel(BaseModel):

    def __init__(self,
                 llm,
                 freeze_llm=True,
                 visual_select_layer=-2,
                 pretrained_pth=None,
                 projector_depth=2,
                 llm_lora=None,
                 visual_encoder_lora=None,
                 use_activation_checkpointing=True,
                 max_position_embeddings=None,
                 hidden_size=512,
                 train_stage='2',
                 enable_long_net=True,
                 long_net_pth=None,
                 projector_pth = None,
                 perceiver_pth = None,
                 #config swin_longnet
                 use_swin_longnet = True,
                 add_abs_pe_to_longnet_inputs = True,
                 
                 longnet_pe_gate_ratio = 0.1,
                 longnet_pe_dropout_rate = 0.1,
                 
                 fourier_dims = 32,
                 
                 # config for Perceiver Resampler
                 use_perceiver_resampler = True,
                 perceiver_num_latents=64,
                 perceiver_depth=2,
                 perceiver_fourier_dims = 32, 
                 perceiver_pe_gate_ratio = 0.1,
                 perceiver_pe_dropout_rate = 0.1
                 ):
        super().__init__()
        
        self.enable_long_net = enable_long_net
        
        if enable_long_net:
            print('enable long net')
        else:
            print('disable long net')
        
        self.freeze_llm = freeze_llm
        self.freeze_visual_encoder = True
        
        self.use_swin_longnet = use_swin_longnet
        
        if train_stage == '0':
            print_log('train_stage == 0', 'current')
            self.freeze_llm = True
            self.freeze_long_net = True
    
        if train_stage == '1':
            print_log('train_stage == 1', 'current')
            self.freeze_llm = True
            self.freeze_long_net = False

        elif train_stage == '2':
            print_log('train_stage == 2', 'current')
            self.freeze_llm = False #False
            self.freeze_long_net = False #False

        with LoadWoInit():
            if isinstance(llm, dict):
                llm = self._dispatch_lm_model_cfg(llm, max_position_embeddings)

            self.llm = self._build_from_cfg_or_module(llm)
            

        self.encoder_name = "LongNet_{}_layers_{}_dim".format(2, 512)
        
        if not self.use_swin_longnet:
            self.LongNet_encoder = make_longnet_from_name(self.encoder_name,
                                                          enable_gradient_checkpoint= False) # , drop_path_rate=0.3, dropout=0.3, segment_length=1024
        else:
            print('use swin long net')
            from .coords_pe import Coord2Embed
            self.add_abs_pe_to_longnet_inputs = add_abs_pe_to_longnet_inputs
            self.coord2embed_longnet = Coord2Embed(out_dim=hidden_size, 
                                                   fourier_dims=fourier_dims).to(dtype=self.llm.dtype)  # hidden_size == LongNet input C
            self.longnet_pe_gate = nn.Parameter(torch.tensor(longnet_pe_gate_ratio,
                                                             dtype=self.llm.dtype))
            
            self.longnet_pe_dropout = nn.Dropout(p = longnet_pe_dropout_rate)  # optional
            
            self.LongNet_encoder = make_swin_longnet_from_name(self.encoder_name,
                                                               keep_dim_after_merge= True,
                                                               merge_size = 2,
                                                               use_rel_pos_2d= False,
                                                               enable_gradient_checkpoint= False
                                                               )
        self.LongNet_encoder = self.LongNet_encoder.to(self.llm.dtype)

        
        self.llm.config.use_cache = False
        dispatch_modules(self.llm)

        self.projector_depth = projector_depth

        projector_config = ProjectorConfig(
            visual_hidden_size=hidden_size,
            llm_hidden_size=self.llm.config.hidden_size,
            depth=self.projector_depth)        

        self.projector = ProjectorModel(projector_config).to(
            self.llm.dtype)
        
        self.use_perceiver_resampler = use_perceiver_resampler
        if self.use_perceiver_resampler:
            # ---- 自动选择 Qwen3 或 Qwen2 的 Perceiver 实现 ----
            self.perceiver_num_latents = perceiver_num_latents
            self.perceiver_depth = perceiver_depth

            from .coords_pe import Coord2Embed
            self.key_pos_enc = Coord2Embed(out_dim=self.hidden_size, 
                                           fourier_dims=perceiver_fourier_dims).to(dtype=self.llm.dtype)
            self.key_pos_gate = nn.Parameter(torch.tensor(perceiver_pe_gate_ratio
                                             , dtype =self.llm.dtype))
            self.key_pos_dropout = nn.Dropout(p = perceiver_pe_dropout_rate)
            
            qwen_major = _detect_qwen_major_version(self.llm)
            print_log(f'using qwen version{qwen_major}', 'current')
            if qwen_major == 3:
                # Qwen3 分支
                try:
                    from .qwen3_perceiver_resampler import (
                        PerceiverResampler as _PR,
                        init_perceiver_from_llm_auto as _init_pr,
                    )
                    print_log('using qwen3', 'current')
                except Exception as e:
                    raise RuntimeError(
                        "检测到 Qwen3，但未找到 qwen3_perceiver_resampler，请确认文件存在且 transformers 版本满足要求（>=4.51）。"
                    ) from e
            elif qwen_major == 2:
                # Qwen2 分支
                from .qwen2_perceiver_resampler import (
                    PerceiverResampler as _PR,
                    init_perceiver_from_llm_auto as _init_pr,
                )
            else:
                warnings.warn(
                    "未能确定 Qwen 主版本（既不是 qwen3 也不是 qwen2）。将回退到 Qwen2 的 Perceiver 实现。",
                    RuntimeWarning,
                )
                from .qwen2_perceiver_resampler import (
                    PerceiverResampler as _PR,
                    init_perceiver_from_llm_auto as _init_pr,
                )

            # 构建并初始化 Perceiver
            self.perceiver = _PR(
                self.llm,
                num_latents=self.perceiver_num_latents,
                depth=self.perceiver_depth,
            ).to(self.llm.dtype)

            _init_pr(
                perceiver=self.perceiver,
                llm=self.llm,
                ckpt_hint=getattr(self.llm.config, "_name_or_path", None),
                init_from_layers=self.perceiver.depth,
                layer_offset=0,
                allow_download=False,
            )
            

        if self.freeze_llm:
            print('freeze_llm')
            self.llm.requires_grad_(False)
             
        if self.freeze_long_net:
            print('freeze_long_net')
            self.LongNet_encoder.requires_grad_(False)
        

        if use_activation_checkpointing:
            # For backward compatibility
            if hasattr(self.llm, 'enable_input_require_grads'):
                self.llm.enable_input_require_grads()
            else:
                self.llm.get_input_embeddings().register_forward_hook(
                    make_inputs_require_grad)
                
            if self.use_perceiver_resampler:
                self.perceiver.enable_input_require_grads()
            
            self.projector.enable_input_require_grads()
            
            # self.LongNet_encoder.enable_input_require_grads()

            # enable gradient (activation) checkpointing for memory efficiency
            self.gradient_checkpointing_enable()

        self.use_llm_lora = llm_lora is not None
        self.use_visual_encoder_lora = None
        if self.use_llm_lora:
            print_log(f"Building lora {llm_lora.__str__}", "current")
            self._prepare_llm_for_lora(llm_lora, use_activation_checkpointing)
            
        # ── 2) Load projector + LongNet from safetensors ────────────────────────
        if long_net_pth is not None:
            print_log(f"Loading LongNet from {long_net_pth}", "current")
            ln_sd = load_file(long_net_pth, device="cpu")
            self.LongNet_encoder.load_state_dict(ln_sd, strict=False)
            self.LongNet_encoder.to(self.llm.dtype)

        if projector_pth is not None:
            print_log(f"Loading projector from {projector_pth}", "current")
            proj_sd = load_file(projector_pth, device="cpu")
            self.projector.load_state_dict(proj_sd, strict=False)
            self.projector.to(self.llm.dtype)
        
        if perceiver_pth is not None and self.use_perceiver_resampler:
            print_log(f'Loading perceiver from {perceiver_pth}", "current ')
            perceiver_sd = load_file(perceiver_pth, device="cpu")
            self.projector.load_state_dict(perceiver_sd, strict=False)
            self.projector.to(self.llm.dtype)

        # ── 3) Optionally load a classic float checkpoint and filter mismatches ──
        if pretrained_pth is not None:
            sd = guess_load_checkpoint(pretrained_pth)
            model_sd = self.state_dict()
            filtered = {
                k: v for k, v in sd.items()
                if k in model_sd and model_sd[k].shape == v.shape
            }
            missing, unexpected = self.load_state_dict(filtered, strict=False)
            print_log(f"Loaded float ckpt from {pretrained_pth}", "current")
            print_log(f"  missing:   {missing}", "current")
            print_log(f"  unexpected:{unexpected}", "current")
        
        self.visual_select_layer = visual_select_layer

        self._is_init = True

        self.is_first_iter = True

    def _parse_lora_config(self, lora_config):
        if isinstance(lora_config, dict) or isinstance(
                lora_config, Config) or isinstance(lora_config, ConfigDict):
            lora_config = BUILDER.build(lora_config)
        return lora_config

    def _prepare_llm_for_lora(self,
                              lora_config,
                              use_activation_checkpointing=True):
        lora_config = self._parse_lora_config(lora_config)
        self.llm = prepare_model_for_kbit_training(
            self.llm, use_activation_checkpointing)
        if lora_config.target_modules is None:
            modules = find_all_linear_names(self.llm)
            lora_config.target_modules = modules
        self.llm = get_peft_model(self.llm, lora_config)

    def _prepare_visual_encoder_for_lora(self,
                                         lora_config,
                                         use_activation_checkpointing=True):
        lora_config = self._parse_lora_config(lora_config)
        if lora_config.target_modules is None:
            modules = find_all_linear_names(self.visual_encoder)
            lora_config.target_modules = modules
        self.visual_encoder = get_peft_model(self.visual_encoder, lora_config)

    def gradient_checkpointing_enable(self):
        self.activation_checkpointing_enable()

    def activation_checkpointing_enable(self):
        self.llm.gradient_checkpointing_enable()
        # self.visual_encoder.gradient_checkpointing_enable()
        self.projector.gradient_checkpointing_enable()
        
        if self.use_perceiver_resampler:
            self.perceiver.enable_input_require_grads()

    def gradient_checkpointing_disable(self):
        self.activation_checkpointing_disable()

    def activation_checkpointing_disable(self):
        self.llm.gradient_checkpointing_disable()
        # self.visual_encoder.gradient_checkpointing_disable()
        self.projector.gradient_checkpointing_disable()
        if self.use_perceiver_resampler:
            self.perceiver.disable_gradient_checkpointing()

    def init_weights(self):
        pass

    def state_dict(self, *args, **kwargs):
        state_dict = super().state_dict(*args, **kwargs)
        to_return = OrderedDict()
        # Step 1. visual_encoder
        if self.use_visual_encoder_lora:
            to_return.update(
                get_peft_model_state_dict(
                    self.visual_encoder, state_dict=state_dict))
        elif not self.freeze_visual_encoder:
            to_return.update({
                k: v
                for k, v in state_dict.items() if 'visual_encoder.' in k
            })
        # Step 2. LLM
        if self.use_llm_lora:
            to_return.update(
                get_peft_model_state_dict(self.llm, state_dict=state_dict))
        
        elif not self.freeze_llm:
            to_return.update(
                {k: v
                 for k, v in state_dict.items() if 'llm.' in k})
        # Step 3. Projector
        to_return.update(
            {k: v
             for k, v in state_dict.items() if 'projector.' in k})
        
        # Step 4. LongNet_encoder
        to_return.update(
            {k: v
             for k, v in state_dict.items() if 'LongNet_encoder.' in k})
        
            # Step 5. Perceiver Resampler (unchanged)
        if getattr(self, 'use_perceiver_resampler', False) and getattr(self, 'perceiver', None) is not None:
            to_return.update({k: v for k, v in state_dict.items() if 'perceiver.' in k})

        # Step 6. NEW — Positional encoders & gates
        # 6a) LongNet input PE
        if hasattr(self, 'coord2embed_longnet'):
            to_return.update({k: v for k, v in state_dict.items() if 'coord2embed_longnet.' in k})
        
        if 'longnet_pe_gate' in state_dict:
            to_return['longnet_pe_gate'] = state_dict['longnet_pe_gate']

        # 6b) Perceiver key-side PE (only if you added it)
        if hasattr(self, 'key_pos_enc'):
            to_return.update({k: v for k, v in state_dict.items() if 'key_pos_enc.' in k})
        if 'key_pos_gate' in state_dict:
            to_return['key_pos_gate'] = state_dict['key_pos_gate']
        
        return to_return

    @staticmethod
    def _prepare_for_long_context_training(cfg, llm_cfg,
                                           max_position_embeddings):

        orig_rope_scaling = getattr(llm_cfg, 'rope_scaling', None)
        if orig_rope_scaling is None:
            orig_rope_scaling = {'factor': 1}

        orig_rope_scaling_factor = orig_rope_scaling[
            'factor'] if 'factor' in orig_rope_scaling.keys() else 1
        orig_ctx_len = getattr(llm_cfg, 'max_position_embeddings', None)
        if orig_ctx_len:
            orig_ctx_len *= orig_rope_scaling_factor
            if max_position_embeddings > orig_ctx_len:
                scaling_factor = float(
                    math.ceil(max_position_embeddings / orig_ctx_len))
                llm_cfg.rope_scaling = {
                    'type': 'linear',
                    'factor': scaling_factor
                }

        # hardcode for internlm2
        llm_cfg.attn_implementation = 'flash_attention_2'
        cfg.config = llm_cfg

        return cfg, llm_cfg

    @staticmethod
    def _prepare_for_flash_attn(cfg, llm_cfg):
        cls_name = type(llm_cfg).__name__
        SUPPORT_SDPA_ATTN = ('LlamaConfig', 'GemmaConfig', 'MistralConfig',
                             'MixtralConfig', 'Qwen2Config', 'Qwen2MoeConfig',
                             'Starcoder2Config', 'Starcoder2Config',
                             'Phi3Config')
        SUPPORT_FLASH_ATTN2 = ('InternLM2Config', 'LlamaConfig', 'GemmaConfig',
                               'MistralConfig', 'MixtralConfig', 'Qwen2Config',
                               'Qwen2MoeConfig', 'Starcoder2Config',
                               'Starcoder2Config', 'Phi3Config')

        torch_dtype = torch.bfloat16 if (
            torch.cuda.is_available() and torch.cuda.is_bf16_supported()) \
            else torch.float16

        if getattr(cfg, 'attn_implementation', None) is not None:
            # Flash Attention 2.0 only supports torch.float16 and
            # torch.bfloat16 dtypes
            if cfg.attn_implementation == 'flash_attention_2':
                cfg.torch_dtype = torch_dtype
        elif SUPPORT_FLASH2 and cls_name in SUPPORT_FLASH_ATTN2:
            cfg.torch_dtype = torch_dtype
            cfg.attn_implementation = 'flash_attention_2'
        elif SUPPORT_FLASH1 and cls_name in SUPPORT_SDPA_ATTN:
            cfg.attn_implementation = 'sdpa'

        return cfg, llm_cfg

    @staticmethod
    def _prepare_for_qlora_zero3(cfg):
        if (not is_deepspeed_zero3_enabled()) or (not hasattr(
                cfg, 'quantization_config')):
            return cfg

        torch_dtype = torch.bfloat16 if (
            torch.cuda.is_available() and torch.cuda.is_bf16_supported()) \
            else torch.float16

        cfg.torch_dtype = torch_dtype
        quantization_config = cfg.quantization_config
        quantization_config.bnb_4bit_compute_dtype = torch_dtype
        quantization_config.bnb_4bit_quant_storage = torch_dtype

        return cfg

    def _dispatch_lm_model_cfg(self, cfg, max_position_embeddings=None):
        cfg = self._prepare_for_qlora_zero3(cfg)
        pretrained_model_name_or_path = cfg.pretrained_model_name_or_path
        llm_cfg = AutoConfig.from_pretrained(
            pretrained_model_name_or_path, trust_remote_code=True)
        
        cfg, llm_cfg = self._prepare_for_flash_attn(cfg, llm_cfg)
        if max_position_embeddings is not None:
            cfg, llm_cfg = self._prepare_for_long_context_training(
                cfg, llm_cfg, max_position_embeddings)
        return cfg

    def _build_from_cfg_or_module(self, cfg_or_mod):
        if isinstance(cfg_or_mod, nn.Module):
            return cfg_or_mod
        elif isinstance(cfg_or_mod, dict):
            traverse_dict(cfg_or_mod)
            return BUILDER.build(cfg_or_mod)
        else:
            raise NotImplementedError

    def forward(self, data, data_samples=None, mode='loss'):
        if self.is_first_iter:
            # hardcode for qlora DeepSpeed ZeRO3, put buffers and QuantState to
            # device
            # Only required in `LLaVAModel` .
            # We do not need this in `SupervisedFinetune` .
            self.to(data['input_ids'].device)
            self.is_first_iter = False
        
        # data_dict['pixel_values']=[[pixel_values of img1], [pixel_values of img2], ...]
        coords_v = None
        if 'pixel_values' in data:
            feat_to_proj = data['pixel_values'].to(self.llm.dtype) # torch.Size([1, img_num, 512])
            if self.enable_long_net:
                if not self.use_swin_longnet:
                    # long_net_output = self.LongNet_encoder(src_tokens=None, token_embeddings=feat_to_proj.permute(1, 0, 2))["encoder_out"] # shape: (img_num, 1, 1024)
                    long_net_output = self.LongNet_encoder(src_tokens=None, token_embeddings=feat_to_proj)["encoder_out"]
                elif self.add_abs_pe_to_longnet_inputs and 'coords' in data:
                    # coords_shape = data['coords'].shape
                    # print_log(f'using swin long net with coords shape {coords_shape}', 'current')
                    pe = self.coord2embed_longnet(data['coords'].to(feat_to_proj.dtype)).to(feat_to_proj.dtype)
                    feat_to_proj = feat_to_proj + self.longnet_pe_dropout(self.longnet_pe_gate * pe)
                    long_net_output = self.LongNet_encoder(src_tokens=None, token_embeddings=feat_to_proj, 
                                                           coords=data['coords'].to(self.llm.dtype))
                    long_net_output, coords_v = long_net_output["encoder_out"], long_net_output['coords']
                elif 'coords' in data:
                    # feat_to_proj = feat_to_proj + self.longnet_pe_dropout(self.longnet_pe_gate * pe)
                    long_net_output = self.LongNet_encoder(src_tokens=None,
                                                           token_embeddings=feat_to_proj, 
                                                           coords=data['coords'].to(self.llm.dtype))
                    long_net_output, coords_v = long_net_output["encoder_out"], long_net_output['coords']
                else:
                    long_net_output = self.LongNet_encoder(src_tokens=None, token_embeddings=feat_to_proj)["encoder_out"]
                # wl - output shape (img_num, 1, 512)
                feat_to_proj = long_net_output # permuted shape: [1, img_num, 512]

            pixel_values = self.projector(feat_to_proj.to(self.llm.dtype)) # output shape [1, patch_num, 3584]
            if self.use_perceiver_resampler and 'input_ids' in data:
                
                # do this only here to void copy embedding layer to preceiver
                text_embeddings = self.llm.get_input_embeddings()(
                    data["input_ids"].clamp(min=0)
                ).to(self.llm.dtype).detach()
                if coords_v is not None:
                    kpe = self.key_pos_enc(coords_v.to(pixel_values.device)).to(pixel_values.dtype)
                    pixel_values = pixel_values + self.key_pos_dropout(self.key_pos_gate * kpe)
                compressed = self.perceiver(
                    # input_ids = data["input_ids"],
                    text_embeddings=text_embeddings,
                    attention_mask=data.get("attention_mask", None),
                    visual_tokens=pixel_values,
                )
                data["pixel_values"] = compressed
            else:
                data['pixel_values'] = pixel_values # shape: [1, patch_num, 3584]   # shape: [1, 576, 4096]
            
            # remove coords
            data.pop('coords', None)
            
            data = prepare_inputs_labels_for_multimodal(llm=self.llm, **data)

        if mode == 'loss':
            return self.compute_loss(data, data_samples)
        elif mode == 'predict':
            return self.predict(data, data_samples)
        elif mode == 'tensor':
            return self._forward(data, data_samples)
        else:
            raise NotImplementedError

    def _forward(self, data, data_samples=None):

        outputs = self.llm(**data)

        return outputs

    def predict(self, data, data_samples=None):
        outputs = self.llm(**data)
        logits_dict = [{'logits': logits} for logits in outputs.logits]
        return logits_dict

    # def compute_loss(self, data, data_samples=None):
    #     outputs = self.llm(**data)
    #     # outputs.logits.shape (1, 1094, 152064) for Qwen
    #     loss_dict = {'loss': outputs.loss}
    #     return loss_dict
        # === MOD: token-averaged, globally weighted loss (robust to variable lengths)
        

    def compute_loss(self, data, data_samples=None):
        # 1) 若无 labels，退回 HF 默认
        if 'labels' not in data:
            outputs = self.llm(**data)
            return {'loss': outputs.loss}

        labels = data['labels']  # [B, T]
        # 不把 labels 交给 HF，避免其先做 per-device mean
        model_inputs = {k: v for k, v in data.items() if k != 'labels'}

        outputs = self.llm(**model_inputs, use_cache=False)
        logits = outputs.logits  # [B, T, V]

        # 2) CausalLM 对齐
        shift_logits = logits[:, :-1, :].contiguous()
        shift_labels = labels[:, 1:].contiguous()

        # 3) 本卡有效 token 数（忽略 -100）
        n_tok_local = (shift_labels != -100).sum().to(device=logits.device, dtype=torch.long)

        # 4) 分子：sum over tokens（用 FP32 计算更稳）
        loss_sum_local = F.cross_entropy(
            shift_logits.float().view(-1, shift_logits.size(-1)),
            shift_labels.view(-1),
            ignore_index=-100,
            reduction='sum'
        )

        # 5) 计算全局分母；不要让反传穿过 collective（用 no_grad + clone）
        world_size = 1
        n_tok_global = n_tok_local
        if dist.is_available() and dist.is_initialized():
            world_size = dist.get_world_size()
            with torch.no_grad():
                n_tok_global = n_tok_local.clone()
                dist.all_reduce(n_tok_global, op=dist.ReduceOp.SUM)

        denom = n_tok_global.clamp_min(1).to(loss_sum_local.dtype)

        # 6) 构造最终 loss：
        #    用“本卡分子 / 全局分母”，再乘 world_size 抵消 DDP 的梯度平均，
        #    这样反向后的等效梯度就是“全局 token 平均”的梯度。
        loss = (loss_sum_local / denom) * float(world_size)

        # 7) 记录指标：把 ntok 作为张量返回，避免 parse_losses 报错
        ntok_tensor = denom.detach()  # float 标量张量即可

        return {
            'loss': loss,
            'ntok': ntok_tensor
        }


    def __getattr__(self, name: str):
        try:
            return super().__getattr__(name)
        except AttributeError:
            return getattr(self.llm, name)

    def to_hf(self,
              cfg,
              save_dir,
              fp32=False,
              save_pretrained_kwargs={},
              save_format='xtuner',
              **kwargs):
        if save_format == 'xtuner':
            self.to_xtuner_llava(cfg, save_dir, fp32, save_pretrained_kwargs)
        elif save_format == 'huggingface':
            self.to_huggingface_llava(cfg, save_dir, fp32,
                                      save_pretrained_kwargs)
        elif save_format == 'official':
            self.to_official_llava(cfg, save_dir, fp32, save_pretrained_kwargs)
        else:
            raise NotImplementedError

    def to_xtuner_llava(self,
                        cfg,
                        save_dir,
                        fp32=False,
                        save_pretrained_kwargs={}):
        # LLM
        self.llm.config.use_cache = True
        if not fp32:
            print_log('Convert LLM to float16', 'current')
            self.llm.half()
        if self.use_llm_lora:
            llm_path = osp.join(save_dir, 'llm_adapter')
            print_log(f'Saving LLM adapter to {llm_path}', 'current')
            self.llm.save_pretrained(llm_path, **save_pretrained_kwargs)
        elif not self.freeze_llm:
            llm_path = save_dir
            print_log(f'Saving LLM tokenizer to {llm_path}', 'current')
            tokenizer = BUILDER.build(cfg.tokenizer)
            tokenizer.save_pretrained(llm_path, **save_pretrained_kwargs)
            print_log(f'Saving LLM to {llm_path}', 'current')
            self.llm.save_pretrained(llm_path, **save_pretrained_kwargs)
        self.llm.config.use_cache = False

        # Visual Encoder
        if self.use_visual_encoder_lora:
            visual_encoder_path = osp.join(save_dir, 'visual_encoder_adapter')
            print_log(
                f'Saving visual_encoder adapter to {visual_encoder_path}',
                'current')
            self.visual_encoder.save_pretrained(visual_encoder_path,
                                                **save_pretrained_kwargs)
        elif not self.freeze_visual_encoder:
            visual_encoder_path = osp.join(save_dir, 'visual_encoder')
            print_log(
                'Saving visual_encoder image_processor to'
                f'{visual_encoder_path}', 'current')
            image_processor = BUILDER.build(cfg.image_processor)
            image_processor.save_pretrained(visual_encoder_path,
                                            **save_pretrained_kwargs)
            print_log(f'Saving visual_encoder to {visual_encoder_path}',
                      'current')
            self.visual_encoder.save_pretrained(visual_encoder_path,
                                                **save_pretrained_kwargs)

        # Projector
        projector_path = osp.join(save_dir, 'projector')
        print_log(f'Saving projector to {projector_path}', 'current')
        # self.projector.save_pretrained(projector_path,
        #                                **save_pretrained_kwargs)
        os.makedirs(projector_path, exist_ok=True)
        output_path = os.path.join(projector_path, 'projector.safetensors')
        save_file(self.projector.state_dict(), output_path)
        
        if self.use_perceiver_resampler:
            
            perceiver_path = osp.join(save_dir, "perceiver")
            print_log(f'Saving LongNet_encoder to {perceiver_path}', 'current')
            os.makedirs(perceiver_path, exist_ok=True)
            perceiver_output_path = os.path.join(perceiver_path, 'perceiver.safetensors')
            save_file(self.perceiver.state_dict(), perceiver_output_path)

        # LongNet_encoder
        if self.LongNet_encoder is not None:
            LongNet_encoder_path = osp.join(save_dir, 'LongNet_encoder')
            print_log(f'Saving LongNet_encoder to {LongNet_encoder_path}', 'current')
            # Ensure the target directory exists
            os.makedirs(LongNet_encoder_path, exist_ok=True)

            # Define the full path for the weights file
            output_path = osp.join(LongNet_encoder_path, 'longnet_encoder.safetensors')

            # Save the state dictionary using safetensors
            save_file(self.LongNet_encoder.state_dict(), output_path)
        
    
       


    def to_huggingface_llava(self,
                             cfg,
                             save_dir,
                             fp32=False,
                             save_pretrained_kwargs={}):

        LLM_MAPPING = {
            'model': 'language_model.model',
            'lm_head': 'language_model.lm_head',
        }
        VIT_MAPPING = {
            'vision_model': 'vision_tower.vision_model',
        }
        PROJECTOR_MAPPING = {
            'model.0': 'multi_modal_projector.linear_1',
            'model.2': 'multi_modal_projector.linear_2',
        }
        LONGNET_MAPPING = {
            'layers.0': 'LongNet_encoder.layers.0',
            'layers.1': 'LongNet_encoder.layers.1',
            'layer_norm': 'LongNet_encoder.layer_norm'
        }

        assert getattr(self.llm, 'hf_quantizer', None) is None, \
            'This conversion format does not support quantized LLM.'

        # get state_dict
        llm = self.llm
        if self.use_llm_lora:
            llm = self.llm.merge_and_unload()
        llm.config.use_cache = True
        if not fp32:
            print_log('Convert LLM to float16', 'current')
            llm.half()

        assert isinstance(llm, LlamaForCausalLM), \
            'This conversion format only supports LlamaForCausalLM.'
        llm_state_dict = llm.state_dict()
        llm_state_dict = convert_state_dict_to_hf(llm_state_dict, LLM_MAPPING)

        need_visual_encoder = (not self.freeze_visual_encoder
                               or self.use_visual_encoder_lora)
        visual_encoder = self.visual_encoder
        if self.use_visual_encoder_lora:
            visual_encoder = self.visual_encoder.merge_and_unload()
        assert isinstance(visual_encoder, CLIPVisionModel),\
            'This conversion format only supports CLIPVisionModel.'
        if need_visual_encoder:
            visual_encoder_state_dict = visual_encoder.state_dict()
            visual_encoder_state_dict = convert_state_dict_to_hf(
                visual_encoder_state_dict, VIT_MAPPING)
        else:
            visual_encoder_state_dict = {}

        projector_state_dict = self.projector.state_dict()
        projector_state_dict = convert_state_dict_to_hf(
            projector_state_dict, PROJECTOR_MAPPING)

        LongNet_encoder_state_dict = self.LongNet_encoder.state_dict()
        LongNet_encoder_state_dict = convert_state_dict_to_hf(
            LongNet_encoder_state_dict, LONGNET_MAPPING)

        state_dict = {
            **projector_state_dict,
            **llm_state_dict,
            **visual_encoder_state_dict,
            **LongNet_encoder_state_dict
        }

        # init model
        text_config = llm.config
        vision_config = visual_encoder.config
        config = LlavaConfig(
            text_config=text_config,
            vision_config=vision_config,
            attn_implementation='eager')

        with init_empty_weights():
            with warnings.catch_warnings():
                warnings.filterwarnings(
                    'ignore', message='.*non-meta.*', category=UserWarning)
                model = LlavaForConditionalGeneration(config)
        model.load_state_dict(state_dict, strict=True, assign=True)

        # processor
        cfg.tokenizer.type = LlamaTokenizerFast.from_pretrained
        tokenizer = BUILDER.build(cfg.tokenizer)

        tokenizer.add_tokens(
            AddedToken(DEFAULT_IMAGE_TOKEN, special=True, normalized=False),
            special_tokens=True)
        tokenizer.add_special_tokens({'pad_token': '<pad>'})

        image_processor = BUILDER.build(cfg.image_processor)
        assert isinstance(image_processor, CLIPImageProcessor),\
            'This conversion format only supports CLIPImageProcessor.'

        processor = LlavaProcessor(
            tokenizer=tokenizer, image_processor=image_processor)

        # Pad to 64 for performance reasons
        pad_shape = 64

        pre_expansion_embeddings = \
            model.language_model.model.embed_tokens.weight.data
        mu = torch.mean(pre_expansion_embeddings, dim=0).float()
        n = pre_expansion_embeddings.size()[0]
        sigma = ((pre_expansion_embeddings - mu).T
                 @ (pre_expansion_embeddings - mu)) / n
        dist = torch.distributions.multivariate_normal.MultivariateNormal(
            mu, covariance_matrix=1e-5 * sigma)

        # We add an image token so we need to resize the model
        ori_vocab_size = config.text_config.vocab_size
        tokenizer_vocab_size = tokenizer.encode('<pad>')[-1]
        added_token = tokenizer_vocab_size - ori_vocab_size

        if added_token > 0:
            model.resize_token_embeddings(ori_vocab_size + added_token,
                                          pad_shape)
            model.language_model.model.embed_tokens.weight.data[
                ori_vocab_size:] = torch.stack(
                    tuple(
                        dist.sample()
                        for _ in range(model.language_model.model.embed_tokens.
                                       weight.data[ori_vocab_size:].shape[0])),
                    dim=0,
                )
            model.language_model.lm_head.weight.data[
                ori_vocab_size:] = torch.stack(
                    tuple(dist.sample()
                          for _ in range(model.language_model.lm_head.weight.
                                         data[ori_vocab_size:].shape[0])),
                    dim=0,
                )
        model.config.image_token_index = tokenizer.encode(
            DEFAULT_IMAGE_TOKEN)[-1]
        model.config.pad_token_id = tokenizer.encode('<pad>')[-1]

        # save
        print_log(f'Saving to {save_dir}', 'current')
        model.save_pretrained(save_dir, **save_pretrained_kwargs)
        processor.save_pretrained(save_dir, **save_pretrained_kwargs)

    def to_official_llava(self,
                          cfg,
                          save_dir,
                          fp32=False,
                          save_pretrained_kwargs={}):

        VIT_MAPPING = {
            'vision_model': 'model.vision_tower.vision_tower.vision_model',
        }
        PROJECTOR_MAPPING = {
            'model.0': 'model.mm_projector.0',
            'model.2': 'model.mm_projector.2',
        }
        LONGNET_MAPPING = {
            'layers.0': 'LongNet_encoder.layers.0',
            'layers.1': 'LongNet_encoder.layers.1',
            'layer_norm': 'LongNet_encoder.layer_norm'
        }

        try:
            from llava.model import LlavaConfig, LlavaLlamaForCausalLM
        except ImportError:
            raise ImportError(
                'Please install llava with '
                '`pip install git+https://github.com/haotian-liu/LLaVA.git '
                '--no-deps`.')

        assert getattr(self.llm, 'hf_quantizer', None) is None, \
            'This conversion format does not support quantized LLM.'

        # get state_dict
        llm = self.llm
        if self.use_llm_lora:
            llm = self.llm.merge_and_unload()
        llm.config.use_cache = True
        if not fp32:
            print_log('Convert LLM to float16', 'current')
            llm.half()

        assert isinstance(llm, LlamaForCausalLM), \
            'This conversion format only supports LlamaForCausalLM.'
        llm_state_dict = llm.state_dict()

        need_visual_encoder = (not self.freeze_visual_encoder
                               or self.use_visual_encoder_lora)
        visual_encoder = self.visual_encoder
        if self.use_visual_encoder_lora:
            visual_encoder = self.visual_encoder.merge_and_unload()
        assert isinstance(visual_encoder, CLIPVisionModel),\
            'This conversion format only supports CLIPVisionModel.'
        if need_visual_encoder:
            visual_encoder_state_dict = visual_encoder.state_dict()
            visual_encoder_state_dict = convert_state_dict_to_hf(
                visual_encoder_state_dict, VIT_MAPPING)
        else:
            visual_encoder_state_dict = {}

        projector_state_dict = self.projector.state_dict()
        projector_state_dict = convert_state_dict_to_hf(
            projector_state_dict, PROJECTOR_MAPPING)

        LongNet_encoder_state_dict = self.LongNet_encoder.state_dict()
        LongNet_encoder_state_dict = convert_state_dict_to_hf(
            LongNet_encoder_state_dict, LONGNET_MAPPING)

        state_dict = {
            **projector_state_dict,
            **llm_state_dict,
            **visual_encoder_state_dict,
            **LongNet_encoder_state_dict
        }

        # init model
        tokenizer = BUILDER.build(cfg.tokenizer)
        image_processor = BUILDER.build(cfg.image_processor)
        assert isinstance(image_processor, CLIPImageProcessor),\
            'This conversion format only supports CLIPImageProcessor.'

        llava_config_dict = llm.config.__dict__.copy()
        llava_config_dict.update(
            dict(
                image_aspect_ratio='pad',
                mm_hidden_size=visual_encoder.config.hidden_size,
                mm_projector_type=f'mlp{self.projector_depth}x_gelu',
                mm_use_im_patch_token=False,
                mm_use_im_start_end=False,
                mm_vision_select_feature='patch',
                mm_vision_select_layer=self.visual_select_layer,
                mm_vision_tower=visual_encoder.config.name_or_path,
                unfreeze_mm_vision_tower=need_visual_encoder,
                model_type='llava',
                use_cache=True,
                use_mm_proj=True))

        llava_config = LlavaConfig(**llava_config_dict)

        with init_empty_weights():
            with warnings.catch_warnings():
                warnings.filterwarnings(
                    'ignore', message='.*non-meta.*', category=UserWarning)
                model = LlavaLlamaForCausalLM(llava_config)

        model.load_state_dict(state_dict, strict=True, assign=True)

        # save
        print_log(f'Saving to {save_dir}', 'current')

        model.save_pretrained(save_dir, **save_pretrained_kwargs)
        image_processor.save_pretrained(save_dir, **save_pretrained_kwargs)
        tokenizer.save_pretrained(save_dir, **save_pretrained_kwargs)