sglangv0.5.2 & support Qwen3-Next-80B-A3B-Instruct

scripts/playground/lora/analyzer.py

@@ -0,0 +1,77 @@
+import glob
+import json
+import os
+import re
+import sys
+
+from tqdm import tqdm
+
+sys.path.append("../../")
+from fix_corrupted_json import clean_json_file
+
+dirpath = "/Users/ying"
+output_file_prefix = "analyzed_log"
+
+time = {}
+tot_time = {}
+size = {}
+
+os.system(f"rm {output_file_prefix}*")
+
+for dirname in glob.glob(os.path.join(dirpath, "trace*")):
+    print(dirname)
+    trace_name = dirname.split("/")[-1]
+    time[trace_name] = {}
+    size[trace_name] = {}
+    total_time = 0
+    for filename in tqdm(glob.glob(os.path.join(dirname, "*.json"))):
+        step_name = filename.split("/")[-1].split(".")[0]
+        step_name = "_".join(step_name.split("_")[1:])
+        if "prefill" not in filename and "decode" not in filename:
+            continue
+
+        match = re.search(r"(prefill|decode)_step_(\d+)\.json", filename)
+        if match:
+            phase = match.group(1)
+            step = match.group(2)
+        else:
+            raise Exception(f"Cannot parse {filename}")
+
+        try:
+            with open(filename, "r") as f:
+                trace = json.load(f)
+        except:
+            clean_json_file(filename, filename)
+            with open(filename, "r") as f:
+                trace = json.load(f)
+
+        for event in trace["traceEvents"]:
+            name = event["name"]
+            if name in ["profile_prefill_step", "profile_decode_step"]:
+                dur = event["dur"] / 1e3
+                time[trace_name][step_name] = dur
+                break
+        total_time += dur
+
+        step = int(step_name.split("_")[-1])
+        with open(os.path.join(dirname, f"size_{step}.json"), "r") as f:
+            size_info = json.load(f)
+        size[trace_name][step_name] = size_info["size"]
+
+    tot_time[trace_name] = total_time
+    time[trace_name] = dict(
+        sorted(time[trace_name].items(), key=lambda x: int(x[0].split("_")[-1]))
+    )
+    size[trace_name] = dict(
+        sorted(size[trace_name].items(), key=lambda x: int(x[0].split("_")[-1]))
+    )
+
+    with open(f"{output_file_prefix}_{trace_name}", "a") as f:
+        for k, v in time[trace_name].items():
+            size_v = size[trace_name][k]
+            print(f"{k:>15}{v:10.2f}\t{size_v}")
+            f.write(f"{k:>15}{v:10.2f}\t{size_v}\n")
+
+with open(f"{output_file_prefix}_total_time", "w") as f:
+    print(tot_time)
+    json.dump(tot_time, f)

scripts/playground/lora/lora_hf_play.py

@@ -0,0 +1,62 @@
+import torch
+from peft import PeftModel
+from transformers import LlamaForCausalLM, LlamaTokenizer
+
+MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
+# ADAPTER = "winddude/wizardLM-LlaMA-LoRA-7B"
+ADAPTER = "/home/ying/test_lora"
+HF_TOKEN = "..."
+
+
+prompt = """
+### Instruction:
+Write a poem about the transformers Python library.
+Mention the word "large language models" in that poem.
+### Response:
+The Transformers are large language models,
+They're used to make predictions on text.
+"""
+
+
+tokenizer = LlamaTokenizer.from_pretrained(MODEL)
+
+base_model = LlamaForCausalLM.from_pretrained(
+    MODEL,
+    device_map="auto",
+    # load_in_8bit=True,
+    torch_dtype=torch.float16,
+    # use_auth_token=HF_TOKEN,
+).cuda()
+
+
+# base model generate
+with torch.no_grad():
+    output_tensors = base_model.generate(
+        input_ids=tokenizer(prompt, return_tensors="pt").input_ids.cuda(),
+        max_new_tokens=32,
+        do_sample=False,
+    )[0]
+
+output = tokenizer.decode(output_tensors, skip_special_tokens=True)
+print("======= base output ========")
+print(output)
+
+
+# peft model generate
+model = PeftModel.from_pretrained(
+    base_model,
+    ADAPTER,
+    torch_dtype=torch.float16,
+    is_trainable=False,
+)
+
+with torch.no_grad():
+    output_tensors = model.generate(
+        input_ids=tokenizer(prompt, return_tensors="pt").input_ids.cuda(),
+        max_new_tokens=32,
+        do_sample=False,
+    )[0]
+
+output = tokenizer.decode(output_tensors, skip_special_tokens=True)
+print("======= peft output ========")
+print(output)

scripts/playground/lora/lora_vllm_play.py

@@ -0,0 +1,30 @@
+from vllm import LLM, SamplingParams
+from vllm.lora.request import LoRARequest
+
+MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
+ADAPTER = "/home/ying/test_lora"
+prompt = """
+### Instruction:
+Write a poem about the transformers Python library.
+Mention the word "large language models" in that poem.
+### Response:
+The Transformers are large language models,
+They're used to make predictions on text.
+"""
+
+
+llm = LLM(model=MODEL, enable_lora=True)
+
+sampling_params = SamplingParams(
+    temperature=0,
+    max_tokens=32,
+)
+
+prompts = [prompt]
+
+outputs = llm.generate(
+    prompts, sampling_params, lora_request=LoRARequest("test_lora", 1, ADAPTER)
+)
+
+print(outputs[0].prompt)
+print(outputs[0].outputs[0].text)