LoRA adapters not saved to `weights`

[ValeKnappich]

In the default config with LoRA, only the base model is being saved to the weights directory. The config indicates that if save_adapter_separately=False, the adapter weights are merged into the main model, but that is not happening. Instead, adapter parameters are simply stripped from the state dict in clean_lora_state_dict.

    save_adapter_separately: Annotated[
        bool,
        Field(
            description="Whether to save LoRA adapters separately before merging into full model weights.",
        ),
    ] = False

[ValeKnappich]

For the record, the lora weights are still stored inside the the trainer checkpoints and can be merged, in case you have to rescue the weights after a long running training job

from pathlib import Path

from torch.distributed.checkpoint.state_dict_loader import _load_state_dict_from_keys
from transformers import AutoConfig, AutoModelForCausalLM

CKPT_DIR = Path("outputs/run_20260131_081816/checkpoints/step_300/trainer")
BASE_MODEL_ID = "Qwen/Qwen3-4B-Instruct-2507"
LORA_RANK = 8
LORA_ALPHA = 16


print("Loading checkpoint")
state_dict = _load_state_dict_from_keys(
    checkpoint_id=str(CKPT_DIR),
)

print("Merging weights")
scaling_factor = LORA_ALPHA / LORA_RANK
merged_state_dict = {}
for k, weight in state_dict["app"]["model"].items():
    if 'lora_A' in k or 'lora_B' in k:
        continue

    lora_a_weights = state_dict["app"]["model"].get(k.replace('weight', 'lora_A.0'))
    lora_b_weights = state_dict["app"]["model"].get(k.replace('weight', 'lora_B.0'))

    if lora_a_weights is not None and lora_b_weights is not None:
        delta_w = (lora_b_weights @ lora_a_weights) * scaling_factor
        merged_weight = weight + delta_w
        merged_state_dict[k] = merged_weight
    else:
        print(f"\tBase weight for {k} has no LoRA adapters.")
        merged_state_dict[k] = weight

print("Loading weights into model")
model = AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(BASE_MODEL_ID))
model.load_state_dict(merged_state_dict)

out_dir = CKPT_DIR.parent / "merged_model"
print(f"Saving model to {out_dir}")
model.save_pretrained(out_dir)

print("Copying tokenizer and config files")
base_model_dir = list(CKPT_DIR.parent.parent.parent.glob("weights/step_*"))[0]
for f in base_model_dir.glob("*"):
    if f.is_file():
        dest = out_dir / f.name
        if not dest.exists():
            print(f"Writing {dest}")
            dest.write_bytes(f.read_bytes())