fix: inject LoRA before loading checkpoint state_dict
move inject_lora() before load_state_dict in _before_wrap so that LoRA adapter weights from a checkpoint are properly restored on training resume. Previously, inject happened after load, causing lora_A/lora_B keys to be silently ignored (strict=False). Co-Authored-By: ccx1324 <2424441089@qq.com>
This commit is contained in:
@@ -110,8 +110,6 @@ class TrainContextBuilder:
|
||||
|
||||
def _before_wrap(m):
|
||||
m = m.to(device=device)
|
||||
if preloaded_state_dict is not None:
|
||||
m.load_state_dict(preloaded_state_dict, strict=False)
|
||||
if cfg.lora is not None:
|
||||
inject_lora(
|
||||
m,
|
||||
@@ -119,6 +117,8 @@ class TrainContextBuilder:
|
||||
alpha=cfg.lora.alpha,
|
||||
target_modules=set(cfg.lora.target_modules),
|
||||
)
|
||||
if preloaded_state_dict is not None:
|
||||
m.load_state_dict(preloaded_state_dict, strict=False)
|
||||
return m
|
||||
|
||||
context = TrainContext(
|
||||
|
||||
Reference in New Issue
Block a user