Views
No views yet
1class LMWithVectorHead(nn.Module):
2 def __init__(self, model_name, lora_config, output_dim=1):
3 super().__init__()
4 backbone = AutoModel.from_pretrained(model_name, device_map='cpu')
5 # backbone.config.use_cache = False
6 self.peft_model = get_peft_model(backbone, lora_config)
7 self.config = backbone.config
8 hidden_size = backbone.config.hidden_size
9 self.vector_head = nn.Linear(hidden_size, output_dim) # 输出维度为 1
10
11 def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
12 """启用梯度检查点,并处理可能的额外参数"""
13 self.peft_model.enable_input_require_grads()
14 if gradient_checkpointing_kwargs is not None:
15 self.peft_model.gradient_checkpointing_enable(**gradient_checkpointing_kwargs)
16 else:
17 self.peft_model.gradient_checkpointing_enable()
18
19 def forward(self, input_ids, attention_mask=None, labels=None):
20 # if hasattr(self.peft_model, "gradient_checkpointing"):
21 # print(f"✅ 梯度检查点已启用 - 当前模式: {self.peft_model.is_gradient_checkpointing}")
22 # else:
23 # print("❌ 梯度检查点未正确初始化")
24 outputs = self.peft_model(
25 input_ids=input_ids,
26 attention_mask=attention_mask,
27 return_dict=True
28 )
29 # 获取最后一个 token 的隐藏状态
30 last_hidden = outputs.last_hidden_state # [B, T, H]
31 cls_hidden = last_hidden[:, -1, :] # [B, H]
32 logits = self.vector_head(cls_hidden) # [B, 1]
33 logits = torch.sigmoid(logits).squeeze(-1) # 添加 sigmoid 并压缩至 [B]
34
35 loss = None
36 if labels is not None:
37 loss_fct = nn.MSELoss() # 使用 MSE 损失
38 loss = loss_fct(logits, labels) # 计算 logits 和 labels 的 MSE
39
40 return CausalLMOutput(
41 loss=loss,
42 logits=logits
43 )1base_model = AutoModel.from_pretrained(args.model_path)
2
3 # 2. 加载训练好的LoRA适配器到基础模型上
4 peft_model = PeftModel.from_pretrained(
5 base_model, # 使用基础模型,而不是model.peft_model
6 args.lora_path,
7 adapter_name="default"
8 )
9
10 # 3. 创建完整模型结构
11 lora_config = LoraConfig(
12 r=args.r,
13 lora_alpha=args.alpha,
14 target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
15 "gate_proj", "up_proj", "down_proj"],
16 lora_dropout=args.lora_dropout,
17 bias="none",
18 )
19 model = LMWithVectorHead(args.model_path, lora_config)
20
21 # 4. 替换为已加载LoRA的模型
22 model.peft_model = peft_model
23
24 # 5. 加载分类头权重
25 state_dict = torch.load(args.vector_head_path, map_location=device)
26 model.vector_head.load_state_dict(state_dict)
27
28 # 6. 激活适配器并移动到设备
29 model.peft_model.set_adapter("default")
30 model = model.to(device)
31
32 # 评估模式
33 model.eval()1/lora
2├── greedy_answer_conf
3│ └── long_qa
4│ └── batchsize16_accumulation8_epochs10_weightdecay0.1_r8_alpha16_loradropout0.0 (training configuration)
5│ ├── best_checkpoints
6│ │ ├── lora_epoch_best/ # Path to LoRA module
7│ │ └── vector_head_epoch_best.pt # Path to Linear Head weights
8│ └── test_losses.json # Test loss for each epoch
9│
10├── hybrid_answer_conf
11│ └── long_qa
12│ ├── batchsize16_accumulation8_epochs10_weightdecay0.1_r8_alpha16_loradropout0.0 (560k samples)
13│ ├── batchsize16_accumulation8_epochs50_weightdecay0.1_r8_alpha16_loradropout0.0_1k_training_samples (1k samples)
14│ └── batchsize16_accumulation8_epochs50_weightdecay0.1_r8_alpha16_loradropout0.0_2k_training_samples (2k samples)
15│
16└── right_answer_conf
17 └── long_qa
18 └── ... # Same format as above
19
20/mlp
21...