NanoVLM-LLaVA

lora_target_modules = ("q_proj", "k_proj", "v_proj", "o_proj")

NanoVLM/
├── configs/                     # 配置文件
│   ├── model_config.py          # 模型架构配置（Vision/Language/Connector 参数）
│   └── training_config.py       # 训练超参数配置（两阶段 lr/bs/epoch 等）
├── src/nanovlm/
│   ├── model/                   # 模型定义
│   │   ├── vision_encoder.py    # SigLIP 视觉编码器封装
│   │   ├── language_model.py    # Qwen2-0.5B 语言模型封装
│   │   ├── connector.py         # MLP 跨模态连接器
│   │   └── nanovlm.py           # 完整 VLM 组装、forward、保存/加载
│   ├── data/                    # 数据处理
│   │   ├── conversation.py      # 对话模板（Qwen ChatML 格式）
│   │   └── dataset.py           # LLaVA 格式数据集加载
│   ├── training/                # 训练系统
│   │   └── trainer.py           # 两阶段训练器（梯度累积/混合精度/断点续训）
│   ├── inference/               # 推理
│   │   └── generator.py         # 自回归生成器 & 交互式对话
│   └── utils/                   # 工具函数
│       └── utils.py             # 随机种子/参数量统计/设备管理
├── scripts/                     # 入口脚本
│   ├── download_models.py       # 下载预训练权重
│   ├── train.py                 # 训练入口
│   └── inference.py             # 推理入口
├── models/                      # 预训练权重（需下载）
├── checkpoints/                 # 训练检查点
├── data/                        # 训练数据
└── requirements.txt

@dataclass
class VisionConfig:
    model_name_or_path: str = "./models/siglip-so400m-patch14-384"
    image_size: int = 384          # 输入分辨率
    freeze: bool = True            # 视觉编码器全程冻结

@dataclass
class LanguageConfig:
    model_name_or_path: str = "./models/Qwen2-0.5B-Instruct"
    hidden_size: int = 1024
    freeze_stage1: bool = True     # 第一阶段冻结
    use_lora_stage2: bool = True   # 第二阶段用 LoRA

@dataclass
class ConnectorConfig:
    connector_type: str = "mlp"
    vision_hidden_size: int = 1152  # SigLIP so400m 输出维度
    llm_hidden_size: int = 1024     # Qwen2-0.5B 输入维度（实际加载后覆盖为 896）
    mlp_hidden_size: int = 2048     # MLP 隐藏层维度
    mlp_depth: int = 2              # MLP 层数
    activation: str = "gelu"

class VisionEncoder(nn.Module):
    def __init__(self, config: VisionConfig):
        super().__init__()
        # 加载预训练 SigLIP 视觉模型
        self.model = AutoModel.from_pretrained(
            config.model_name_or_path,
            dtype=torch.float16,
            local_files_only=True
        )
        self.vision_model = getattr(self.model, "vision_model", self.model)

        # 判断是否有 CLS token（CLIP 系列有，SigLIP 视觉塔没有）
        model_type = getattr(self.vision_config, "model_type", "")
        self.has_cls_token = model_type.startswith("clip")

        # 冻结所有参数
        if config.freeze:
            self.model.eval()
            for param in self.model.parameters():
                param.requires_grad = False

    def forward(self, pixel_values):
        if self.config.freeze:
            with torch.no_grad():
                outputs = self.vision_model(pixel_values=pixel_values, return_dict=True)
        else:
            outputs = self.vision_model(pixel_values=pixel_values, return_dict=True)
        return outputs.last_hidden_state  # (B, num_patches+1, hidden_size)

class Connector(nn.Module):
    """
    输入:  (B, 729, 1152)   # SigLIP patch features
    输出:  (B, 729, 896)    # Qwen2 embedding 空间

    架构: Linear(1152 → 2048) → GELU → Linear(2048 → 896)
    """

    def __init__(self, config: ConnectorConfig):
        super().__init__()
        layers = []
        input_dim = config.vision_hidden_size  # 1152

        for i in range(config.mlp_depth):  # mlp_depth=2
            if i < config.mlp_depth - 1:
                # 隐藏层: 1152 → 2048
                layers.append(nn.Linear(input_dim, config.mlp_hidden_size))
                input_dim = config.mlp_hidden_size
            else:
                # 输出层: 2048 → 896
                layers.append(nn.Linear(input_dim, config.llm_hidden_size))

            if i < config.mlp_depth - 1:
                layers.append(nn.GELU())   # 隐藏层后加 GELU 激活

        self.mlp = nn.Sequential(*layers)
        self._init_weights()  # Xavier 初始化，gain=0.01

    def forward(self, vision_features):
        return self.mlp(vision_features)

class LanguageModelWrapper(nn.Module):
    def __init__(self, config: LanguageConfig):
        super().__init__()
        self.model = AutoModelForCausalLM.from_pretrained(
            config.model_name_or_path,
            dtype=torch.float16,
            trust_remote_code=True,
            local_files_only=True
        )
        self.tokenizer = AutoTokenizer.from_pretrained(
            config.model_name_or_path,
            use_fast=False,
            local_files_only=True
        )
        # 保证有 pad_token
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        # 启用 gradient checkpointing（用计算换显存）
        self.model.gradient_checkpointing_enable()

class NanoVLM(nn.Module):
    def __init__(self, config: NanoVLMConfig):
        super().__init__()
        self.vision_encoder = VisionEncoder(config.vision)
        self.language_model = LanguageModelWrapper(config.language)

        # 用真实加载到的模型维度覆盖手工配置
        self.config.connector.vision_hidden_size = self.vision_encoder.hidden_size
        self.config.connector.llm_hidden_size = self.language_model.hidden_size
        self.connector = Connector(config.connector)

        self._setup_image_token()  # 注册 <image> 特殊 token

def encode_images(self, images: torch.Tensor) -> torch.Tensor:
    # 1. Vision Encoder: (B,3,384,384) → (B, num_patches+1, 1152)
    vision_features = self.vision_encoder(images)

    # 2. 去掉 CLS token（如果有的话），只保留 patch tokens
    if self.vision_encoder.has_cls_token:
        patch_features = vision_features[:, 1:, :]
    else:
        patch_features = vision_features

    # 3. MLP 投影到 LLM 空间: (B, 729, 1152) → (B, 729, 896)
    visual_embeddings = self.connector(patch_features)
    return visual_embeddings

def prepare_inputs_embeds(self, input_ids, visual_embeddings, attention_mask=None):
    # 1. 获取文本 embedding: (B, L) → (B, L, 896)
    embed_tokens = self.language_model.get_input_embeddings()
    text_embeds = embed_tokens(input_ids)

    # 2. 找到所有 <image> token 的位置
    image_token_mask = (input_ids == self.image_token_id)  # (B, L)

    # 3. 逐样本替换 <image> 位置为视觉 embedding
    for b in range(batch_size):
        image_positions = image_token_mask[b].nonzero(as_tuple=True)[0]
        num_image_tokens = len(image_positions)
        num_vis_tokens = visual_embeddings.shape[1]  # 729

        # 处理数量不匹配的情况
        if num_image_tokens < num_vis_tokens:
            visual_embeds_used = visual_embeddings[b, :num_image_tokens, :]  # 截断
        elif num_image_tokens > num_vis_tokens:
            padding = torch.zeros(num_image_tokens - num_vis_tokens, hidden_size, ...)
            visual_embeds_used = torch.cat([visual_embeddings[b], padding], dim=0)
        else:
            visual_embeds_used = visual_embeddings[b]

        # 替换
        text_embeds[b, image_positions] = visual_embeds_used

    return text_embeds, attention_mask

def forward(self, pixel_values=None, input_ids=None,
            attention_mask=None, labels=None, **kwargs):
    # 1. 图像 → 视觉 embedding
    if pixel_values is not None:
        visual_embeddings = self.encode_images(pixel_values)

    # 2. 构建混合 embedding 序列
    if visual_embeddings is not None and input_ids is not None:
        inputs_embeds, attention_mask = self.prepare_inputs_embeds(
            input_ids, visual_embeddings, attention_mask)

        # 3. 送入 LLM（直接传 embedding，不传 input_ids）
        outputs = self.language_model(
            input_ids=None,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            labels=labels,
        )
    else:
        # 纯文本模式
        outputs = self.language_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
        )

    return {"loss": outputs.loss, "logits": outputs.logits}

<|im_start|>system
You are a helpful vision-language assistant.<|im_end|>
<|im_start|>user
<image> × 729
请描述这张图片。<|im_end|>
<|im_start|>assistant
这是一张...的图片。<|im_end|>

# assistant 回复之前的部分（system + user）的 labels 设为 -100
# 只有 assistant 回复部分的 labels 保留真实 token id
assistant_start_marker = "<|im_start|>assistant\n"
assistant_start_tokens = tokenizer.encode(assistant_start_marker)

# 在 input_ids 中搜索 assistant 起始位置
for i in range(len(input_ids) - len(assistant_start_tokens) + 1):
    if input_ids[i:i+len(assistant_start_tokens)] == assistant_start_tokens:
        ass_start_pos = i + len(assistant_start_tokens)
        break

# 将 assistant 之前的所有 token 的 label 设为 -100
for i in range(ass_start_pos):
    labels[i] = -100

# NanoVLM.set_stage("stage1") 触发:

# VisionEncoder: 冻结（构造时已设置 freeze=True）
# LanguageModel:
self.model.eval()
for param in self.model.parameters():
    param.requires_grad = False

# Connector: 保持可训练（默认 requires_grad=True）

@dataclass
class Stage1Config:
    num_epochs: int = 1
    per_device_batch_size: int = 1
    gradient_accumulation_steps: int = 8   # 等效 batch_size = 8
    learning_rate: float = 2e-3            # 高学习率，快速对齐
    warmup_ratio: float = 0.03             # 前 3% 步数预热
    weight_decay: float = 0.0              # 不对齐做正则化
    use_fp16: bool = True                  # 混合精度节省显存
    max_seq_length: int = 2048

def train_stage1(self, stage1_config):
    # 1. 设置阶段
    self.model.set_stage("stage1")

    # 2. 创建 DataLoader, Optimizer (AdamW), Scheduler (warmup + cosine)
    optimizer, scheduler = self._get_optimizer_and_scheduler(stage1_config, total_steps)
    scaler = torch.amp.GradScaler('cuda', enabled=stage1_config.use_fp16)

    # 3. 训练循环
    for epoch in range(stage1_config.num_epochs):
        for step, batch in enumerate(train_loader):
            # 混合精度前向
            with torch.amp.autocast('cuda', enabled=stage1_config.use_fp16):
                outputs = self.model(pixel_values=pixel_values,
                                     input_ids=input_ids,
                                     attention_mask=attention_mask,
                                     labels=labels)
                loss = outputs["loss"] / gradient_accumulation_steps

            # 梯度累积
            scaler.scale(loss).backward()

            if (step + 1) % gradient_accumulation_steps == 0:
                scaler.unscale_(optimizer)
                grad_norm = clip_grad_norm_(model.parameters(), max_grad_norm=1.0)
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad()
                scheduler.step()   # warmup → cosine 衰减

@dataclass
class Stage2Config:
    learning_rate: float = 2e-4              # 比 Stage1 低一个数量级
    gradient_accumulation_steps: int = 16    # 等效 batch_size = 16
    # LoRA 参数
    lora_r: int = 8                          # 低秩维度
    lora_alpha: int = 16                     # 缩放因子
    lora_dropout: float = 0.05
    lora_target_modules: tuple = ("q_proj", "v_proj", "k_proj", "o_proj")

def _apply_lora(self, stage2_config):
    from peft import LoraConfig, get_peft_model, TaskType

    lora_config = LoraConfig(
        r=8,                    # 低秩矩阵的秩
        lora_alpha=16,          # LoRA 输出的缩放系数
        lora_dropout=0.05,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # 只对注意力投影加 LoRA
        task_type=TaskType.CAUSAL_LM,
    )
    self.model.language_model.model = get_peft_model(
        self.model.language_model.model, lora_config
    )

def save_pretrained(self, save_dir, save_config=True, save_tokenizer=True):
    # 保存 Connector 权重（训练的主要产出）
    torch.save(self.connector.state_dict(), "connector.bin")

    # 保存 LoRA 权重（Stage 2 独有）
    if isinstance(self.language_model.model, PeftModel):
        self.language_model.model.save_pretrained("lora/")

    # 保存 config.json + tokenizer

def generate(self, image, question):
    # 1. 构建对话（system + user + generation_prompt）
    prompt = conversation.apply_chat_template(conversations, add_generation_prompt=True)
    # "<|im_start|>system\n...<|im_end|>\n<|im_start|>user\n<image>×729...<|im_end|>\n<|im_start|>assistant\n"

    # 2. 预处理图像: PIL → pixel_values (1, 3, 384, 384)

    # 3. Tokenize 文本

    # 4. encode_images → visual_embeddings (1, 729, 896)

    # 5. prepare_inputs_embeds → 混合 embedding 序列

    # 6. 自回归生成（KV Cache）
    for _ in range(max_new_tokens):
        if past_key_values is None:
            # 首步：传入完整 embedding 序列
            outputs = lm_model(inputs_embeds=inputs_embeds, use_cache=True)
        else:
            # 后续步骤：只传入最后 1 个 token 的 embedding + KV Cache
            current_embed = embed_tokens(generated[:, -1:])
            outputs = lm_model(inputs_embeds=current_embed,
                               past_key_values=past_key_values, use_cache=True)

        logits = outputs.logits[:, -1, :]
        next_token = sample(logits, temperature, top_k, top_p)
        generated = torch.cat([generated, next_token], dim=-1)

        if next_token == eos_token_id:
            break

    # 7. 解码输出
    response = tokenizer.decode(generated[0, input_len:])