Upload fuse2_qlora_infer.py

fuse2_qlora_infer.py

@@ -0,0 +1,250 @@
+"""Quick local inference for the QLoRA-merged Fuse-2 model.
+
+Loads the model in 4-bit NF4 (fits in 12 GB VRAM), applies the runtime
+patches (expert scale, SwiGLU clamp, router stability), and generates
+text for a given prompt.
+"""
+import os
+import sys
+import gc
+import json
+import time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+sys.stdout.reconfigure(encoding='utf-8', errors='replace')
+sys.stderr.reconfigure(encoding='utf-8', errors='replace')
+
+sys.path.insert(0, "E:/fuse1/mini-deepseek-v4-flash-qlora")
+
+MODEL_PATH = "E:/fuse1/mini-deepseek-v4-flash-qlora"
+DEVICE = "cuda:0"
+SWIGLU_LIMIT = 10.0
+TARGET_STD = 0.025
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.benchmark = True
+
+
+def load_model_4bit():
+    from transformers import AutoConfig, AutoTokenizer
+    from accelerate import init_empty_weights
+    from safetensors import safe_open
+    import bitsandbytes as bnb
+    import fuse2_model_local
+
+    print(f"[load] Config from {MODEL_PATH}...")
+    config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True)
+
+    with init_empty_weights():
+        model = fuse2_model_local.Fuse2ForCausalLM(config)
+
+    with open(f"{MODEL_PATH}/model.safetensors.index.json") as f:
+        index = json.load(f)
+    weight_map = index["weight_map"]
+    shards = sorted(set(weight_map.values()))
+    param_names = set(dict(model.named_parameters()).keys())
+    buffer_names = set(dict(model.named_buffers()).keys())
+
+    tok = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+    if tok.pad_token_id is None:
+        tok.pad_token_id = tok.eos_token_id
+
+    quantize_suffixes = (
+        "q_proj.weight", "k_proj.weight", "v_proj.weight", "o_proj.weight",
+        "gate_proj.weight", "up_proj.weight", "down_proj.weight",
+    )
+
+    def should_quantize(key):
+        return any(key.endswith(s) for s in quantize_suffixes)
+
+    def navigate_to_parent(model, key):
+        parts = key.split(".")
+        obj = model
+        for p in parts[:-1]:
+            if p.isdigit():
+                obj = obj[int(p)]
+            else:
+                obj = getattr(obj, p)
+        return obj, parts[-1]
+
+    def find_linear_owner(model, key):
+        parts = key.split(".")
+        obj = model
+        for p in parts[:-2]:
+            if p.isdigit():
+                obj = obj[int(p)]
+            else:
+                obj = getattr(obj, p)
+        return obj, parts[-2]
+
+    print(f"[load] Loading weights as 4-bit NF4...")
+    total_4bit = 0
+    total_bf16 = 0
+    replaced_linears = {}
+
+    for shard_name in shards:
+        shard_path = os.path.join(MODEL_PATH, shard_name)
+        shard_keys = [k for k, v in weight_map.items() if v == shard_name]
+        print(f"  {shard_name} — {len(shard_keys)} tensors...", flush=True)
+
+        with safe_open(shard_path, framework="pt", device="cpu") as f:
+            for key in shard_keys:
+                if key not in param_names and key not in buffer_names:
+                    continue
+                tensor = f.get_tensor(key)
+
+                if should_quantize(key) and key in param_names:
+                    owner, linear_attr = find_linear_owner(model, key)
+                    linear_id = id(getattr(owner, linear_attr))
+                    if linear_id not in replaced_linears:
+                        old_linear = getattr(owner, linear_attr)
+                        in_f = old_linear.in_features
+                        out_f = old_linear.out_features
+                        has_bias = old_linear.bias is not None
+                        new_linear = bnb.nn.Linear4bit(
+                            in_f, out_f, bias=has_bias,
+                            quant_type="nf4", compute_dtype=torch.bfloat16,
+                            device=DEVICE,
+                        )
+                        t = tensor.to(torch.bfloat16)
+                        new_linear.weight = bnb.nn.Params4bit(
+                            t, requires_grad=False, quant_type="nf4",
+                        ).cuda(0)
+                        if has_bias:
+                            new_linear.bias = None
+                        setattr(owner, linear_attr, new_linear)
+                        replaced_linears[id(new_linear)] = new_linear
+                        del old_linear, new_linear, t
+                        total_4bit += 1
+                else:
+                    parent, param_name = navigate_to_parent(model, key)
+                    t = tensor.to(torch.bfloat16).to(DEVICE)
+                    parent._parameters[param_name] = nn.Parameter(t, requires_grad=False)
+                    del t
+                    total_bf16 += 1
+                del tensor
+                if (total_4bit + total_bf16) % 200 == 0:
+                    gc.collect()
+                    torch.cuda.empty_cache()
+
+        gc.collect()
+        torch.cuda.empty_cache()
+        vram = torch.cuda.memory_allocated(0) / 1e9
+        print(f"    {total_4bit}q+{total_bf16}b, VRAM: {vram:.1f} GB", flush=True)
+
+    # Fix tied embeddings
+    if model.lm_head.weight.device.type == 'meta':
+        model.lm_head.weight = nn.Parameter(
+            model.model.embed_tokens.weight.data.clone(), requires_grad=False
+        )
+
+    # Init coding_gate and coding_norm
+    from fuse2_model_local import Fuse2AugmentedLayer
+    for layer in model.model.layers:
+        if not isinstance(layer, Fuse2AugmentedLayer):
+            continue
+        if hasattr(layer, 'coding_gate') and layer.coding_gate.device.type == 'meta':
+            layer.coding_gate = nn.Parameter(torch.tensor(-2.0, device=DEVICE))
+        if hasattr(layer, 'coding_norm') and layer.coding_norm.weight.device.type == 'meta':
+            layer.coding_norm = nn.RMSNorm(layer.coding_norm.weight.shape[0], eps=1e-6).to(DEVICE)
+
+    # Patch experts (SwiGLU clamp)
+    patched = 0
+    for layer in model.model.layers:
+        if not isinstance(layer, Fuse2AugmentedLayer):
+            continue
+        experts = getattr(layer, "experts", None)
+        if experts is None:
+            continue
+        for expert in experts:
+            gp, up, dp = expert.gate_proj, expert.up_proj, expert.down_proj
+            def make_fwd(g, u, d, lim):
+                def forward(x):
+                    gate = F.silu(g(x))
+                    val = torch.clamp(gate * u(x), -lim, lim)
+                    return d(val)
+                return forward
+            expert.forward = make_fwd(gp, up, dp, SWIGLU_LIMIT)
+            patched += 1
+    print(f"[patch] {patched} experts patched with SwiGLU clamp")
+
+    # Router stability
+    for layer in model.model.layers:
+        if not isinstance(layer, Fuse2AugmentedLayer):
+            continue
+        router = layer.router
+        top_k = router.top_k
+        gate = router.gate
+        def make_stable_fwd(g, tk):
+            def forward(hidden_states):
+                logits = g(hidden_states)
+                scores = F.softplus(logits)
+                scores = torch.clamp(scores, min=1e-6)
+                scores = scores.sqrt()
+                topk_weights, topk_indices = scores.topk(tk, dim=-1)
+                topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-8)
+                return topk_weights, topk_indices, logits
+            return forward
+        router.forward = make_stable_fwd(gate, top_k)
+    print(f"[patch] Router stability applied")
+
+    model.set_coding_enabled(True)
+    model.to(DEVICE)
+    model.eval()
+
+    vram = torch.cuda.memory_allocated(0) / 1e9
+    print(f"[load] Ready, VRAM: {vram:.1f} GB")
+    return model, tok
+
+
+def generate(model, tok, prompt, max_new_tokens=1024, temperature=0.6, repetition_penalty=1.1):
+    """Generate text from a prompt."""
+    # Format as chat
+    messages = [{"role": "user", "content": prompt}]
+    text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    print(f"[gen] Prompt: {prompt[:100]}...")
+    print(f"[gen] Max tokens: {max_new_tokens}, temp: {temperature}")
+    print(f"[gen] Generating...", flush=True)
+
+    inputs = tok(text, return_tensors="pt").to(DEVICE)
+    input_len = inputs["input_ids"].shape[1]
+
+    t0 = time.time()
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            do_sample=temperature > 0,
+            repetition_penalty=repetition_penalty,
+            pad_token_id=tok.pad_token_id,
+            use_cache=True,
+        )
+    t1 = time.time()
+
+    generated = tok.decode(outputs[0][input_len:], skip_special_tokens=True)
+    elapsed = t1 - t0
+    tok_count = outputs.shape[1] - input_len
+    tps = tok_count / elapsed
+
+    print(f"[gen] {tok_count} tokens in {elapsed:.1f}s ({tps:.1f} tok/s)")
+    return generated
+
+
+if __name__ == "__main__":
+    prompt = sys.argv[1] if len(sys.argv) > 1 else "Write a Python fizzbuzz."
+
+    model, tok = load_model_4bit()
+
+    print("\n" + "=" * 60)
+    print("GENERATING")
+    print("=" * 60)
+
+    result = generate(model, tok, prompt, max_new_tokens=2048, temperature=0.6)
+
+    print("\n" + "=" * 60)
+    print("OUTPUT")
+    print("=" * 60)
+    print(result)