File size: 5,964 Bytes
fafdd6c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch",
#     "peft>=0.18.0",
#     "transformers>=5.1.0",
#     "accelerate>=1.6.0",
#     "huggingface_hub",
# ]
# ///
"""
Merge LoRA Adapter and Push to Hub (HF Cloud)

This script is designed to run on Hugging Face Jobs infrastructure.
It loads a base model + LoRA adapter, merges them, and pushes the
standalone merged model to the Hub.

Usage:
    # Submit via Makefile
    make merge-cloud

    # Or submit directly
    python scripts/submit_cloud_training.py \
        --flavor a100x4 --timeout 2h \
        --script scripts/merge_hf_cloud.py
"""

import os

import torch
from huggingface_hub import login
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# ═══════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════

BASE_MODEL = "Qwen/Qwen3-Coder-Next"
ADAPTER_MODEL = "yangluo/univer-api-qwen3-next-coder"
MERGED_MODEL_ID = "yangluo/univer-api-qwen3-next-coder-merged"

# ═══════════════════════════════════════════════════════════════
# MERGE SCRIPT
# ═══════════════════════════════════════════════════════════════


def main():
    # Login to Hugging Face Hub
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        print("Logging in to Hugging Face Hub...")
        login(token=hf_token, add_to_git_credential=False)
        print("  Login successful!")
    else:
        print("ERROR: HF_TOKEN not found in environment!")
        print("Cannot push merged model without authentication.")
        return 1

    print("=" * 60)
    print("    MERGE LoRA ADAPTER & PUSH TO HUB")
    print("=" * 60)
    print(f"Base model:    {BASE_MODEL}")
    print(f"LoRA adapter:  {ADAPTER_MODEL}")
    print(f"Merged output: {MERGED_MODEL_ID}")
    print(f"Dtype:         bfloat16 (full precision, no quantization)")
    print("=" * 60)

    # ─────────────────────────────────────────────────────
    # Load Tokenizer
    # ─────────────────────────────────────────────────────
    print("\n[1/5] Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    print(f"  Tokenizer loaded: {tokenizer.__class__.__name__}")

    # ─────────────────────────────────────────────────────
    # Load Base Model (bf16, no quantization)
    # ─────────────────────────────────────────────────────
    print(f"\n[2/5] Loading base model: {BASE_MODEL}")
    print("  Loading in bfloat16 with device_map='auto'...")

    n_gpus = torch.cuda.device_count()
    print(f"  Detected {n_gpus} GPU(s)")

    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )

    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3
        print(f"  Base model loaded, GPU memory: {allocated:.1f} GB")

    # ─────────────────────────────────────────────────────
    # Load LoRA Adapter
    # ─────────────────────────────────────────────────────
    print(f"\n[3/5] Loading LoRA adapter: {ADAPTER_MODEL}")
    model = PeftModel.from_pretrained(model, ADAPTER_MODEL)
    print("  Adapter loaded successfully")

    # ─────────────────────────────────────────────────────
    # Merge and Unload
    # ─────────────────────────────────────────────────────
    print("\n[4/5] Merging adapter into base model...")
    model = model.merge_and_unload()
    print("  Merge complete!")

    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3
        print(f"  GPU memory after merge: {allocated:.1f} GB")

    # ─────────────────────────────────────────────────────
    # Push to Hub
    # ─────────────────────────────────────────────────────
    print(f"\n[5/5] Pushing merged model to Hub: {MERGED_MODEL_ID}")
    print("  Pushing model (this may take a while for large models)...")
    model.push_to_hub(MERGED_MODEL_ID)
    print("  Model pushed!")

    print("  Pushing tokenizer...")
    tokenizer.push_to_hub(MERGED_MODEL_ID)
    print("  Tokenizer pushed!")

    print("\n" + "=" * 60)
    print("         MERGE COMPLETE!")
    print("=" * 60)
    print(f"Merged model: https://huggingface.co/{MERGED_MODEL_ID}")


if __name__ == "__main__":
    exit(main() or 0)