vllm serve /data/hf/GLM-5-FP8/ --tensor-parallel-size 8 --tool-call-parser glm47 --reasoning-parser glm45 --enable-auto-tool-choice --served-model-name glm-5-fp8

---

import requests
import os
import math
import numpy as np
from transformers import AutoTokenizer

VLLM_URL = "http://localhost:8000/v1/completions"
MODEL_PATH = "/data/hf/GLM-5-FP8/"
PPL_NUM = 5

def get_wikitext(wikitext_path="wiki.test.raw"):
    texts = open(wikitext_path, "r", encoding="utf8").readlines()
    texts = ['' if ttt==' \n' else ttt for ttt in texts]
    return [ttt for ttt in texts if len(ttt) > 0]

def get_perplexity(text: str, model: str = MODEL_PATH):
    """Send text to vLLM server and calculate perplexity from returned logprobs."""
    response = requests.post(VLLM_URL, json={
        "model": "glm-5-fp8",
        "prompt": text,
        "max_tokens": 1,        # we don't need generation, just logprobs of input
        "echo": True,           # return logprobs for prompt tokens too
        "logprobs": 1,          # return top-1 logprob per token (the actual token's logprob)
        "temperature": 0,
    })
    result = response.json()

    logprobs_data = result["choices"][0]["logprobs"]
    token_logprobs = logprobs_data["token_logprobs"]

    valid_logprobs = [lp for lp in token_logprobs if lp is not None]

    # PPL = exp(-1/N * sum(log_probs))
    avg_neg_logprob = -np.mean(valid_logprobs)
    ppl = math.exp(avg_neg_logprob)

    return ppl, valid_logprobs, logprobs_data["tokens"]


def calculate_ppl_for_wikitext():
    """Calculate PPL for first 5 segments of wikitext, each segment has exactly 2048 tokens."""
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
    
    # Get wikitext and concatenate
    texts = get_wikitext()
    full_text = "\n\n".join(texts)
    
    # Tokenize entire text
    encoding = tokenizer(full_text, return_tensors="pt")
    all_input_ids = encoding.input_ids[0]
    seq_len = len(all_input_ids)
    context_len = 2048
    
    print(f"Total tokens: {seq_len}, context length: {context_len}")

    all_ppls = []
    for i in range(PPL_NUM):
        with open("/tmp/ppl_data_idx.txt", "w") as f:
            f.write(str(i))
        begin_loc = i * context_len
        end_loc = min(begin_loc + context_len, seq_len)
        
        if begin_loc >= seq_len:
            break

        segment_ids = all_input_ids[begin_loc:end_loc]
        segment_text = tokenizer.decode(segment_ids, skip_special_tokens=False)

        ppl, logprobs, tokens = get_perplexity(segment_text)
        all_ppls.append(ppl)
        print(f"Segment {i+1}: vllm_tokens={len(segment_ids)}, PPL={ppl:.4f}")
        print(segment_ids)
    
    print(f"\nAverage PPL of 5 segments: {np.mean(all_ppls):.4f}")
    return all_ppls

if __name__ == "__main__":
    calculate_ppl_for_wikitext()

---

import os
import torch
import numpy as np
import time
import logging
from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)


def get_wikitext(wikitext_path="wiki.test.raw"):
    texts = open(wikitext_path, "r", encoding="utf8").readlines()
    texts = ['' if ttt == ' \n' else ttt for ttt in texts]
    return [ttt for ttt in texts if len(ttt) > 0]


def get_perplexity_direct(model, tokenizer, input_ids):
    t0 = time.time()
    with torch.no_grad():
        outputs = model(input_ids=input_ids)
    t1 = time.time()

    if isinstance(outputs, tuple):
        logits = outputs[0]
    else:
        logits = outputs.logits if hasattr(outputs, 'logits') else outputs['logits']
    logger.info(f"  forward time: {t1 - t0:.2f}s, logits shape: {logits.shape}")

    logits_f32 = logits.to(torch.float32)
    probs = torch.softmax(logits_f32[0, :-1], dim=-1)
    labels = input_ids[0, 1:]
    target_probs = probs[torch.arange(len(labels)), labels]
    log2_probs = target_probs.log2().cpu().numpy().tolist()
    ppl = 2 ** (-np.mean(log2_probs))
    return ppl, logits


def calculate_ppl_for_wikitext(model, tokenizer, all_input_ids, logits_dir="logits_output"):
    """Calculate PPL on wikitext segments and save logits to disk."""
    all_input_ids = all_input_ids.to(model.device)
    os.makedirs(logits_dir, exist_ok=True)
    logger.info(f"Logits will be saved to: {os.path.abspath(logits_dir)}")

    seg_len = 2048
    num_segments = 5
    total_tokens = all_input_ids.shape[1]

    ppls = []
    for i in range(num_segments):
        start = i * seg_len
        end = start + seg_len
        if end > total_tokens:
            logger.info(f"Segment {i}: not enough tokens (need {end}, have {total_tokens}), stopping.")
            break
        segment_ids = all_input_ids[:, start:end]
        t0 = time.time()
        ppl, logits = get_perplexity_direct(model, tokenizer, segment_ids)
        t1 = time.time()
        ppls.append(ppl)
        logger.info(f"Segment {i}: PPL = {ppl:.4f}, total time: {t1 - t0:.2f}s")

        # Save logits and input_ids to disk
        save_path = os.path.join(logits_dir, f"segment_{i}.pt")
        torch.save({"logits": logits.cpu(), "input_ids": segment_ids.cpu()}, save_path)
        logger.info(f"  saved logits {logits.shape} + input_ids {segment_ids.shape} to {save_path}")

    if ppls:
        avg_ppl = np.mean(ppls)
        logger.info(f"Average PPL over {len(ppls)} segments: {avg_ppl:.4f}")
    else:
        logger.info("No segments were evaluated.")


if __name__ == "__main__":
    # Load dataset first — fail fast before expensive model loading
    t0 = time.time()
    texts = get_wikitext()
    full_text = "".join(texts)
    t1 = time.time()
    logger.info(f"Loaded wikitext: {len(texts)} lines, {len(full_text)} chars, time: {t1 - t0:.2f}s")

    model_path = "/data/hf/GLM-5-FP8/"

    t0 = time.time()
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    t1 = time.time()
    logger.info(f"Tokenizer loaded, time: {t1 - t0:.2f}s")

    t0 = time.time()
    encoded = tokenizer(full_text, return_tensors="pt")
    all_input_ids = encoded["input_ids"]
    t1 = time.time()
    total_tokens = all_input_ids.shape[1]
    logger.info(f"Tokenized wikitext: {total_tokens} tokens, input_ids shape: {all_input_ids.shape}, time: {t1 - t0:.2f}s")
    if total_tokens < 2048:
        raise RuntimeError(f"Not enough tokens ({total_tokens}) for even one 2048-token segment")

    t0 = time.time()
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=False)
    config._attn_implementation = "eager" # without setting this, the ppl is around 20k,

    from accelerate import dispatch_model
    max_memory = {i: "144GiB" for i in range(torch.cuda.device_count())}
    model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, config=config, device_map="cpu")

    device_map = {"model.embed_tokens": 0}
    layers_per_gpu = [10, 10, 10, 10, 10, 10, 10, 8]
    layer_idx = 0
    for gpu_id, n_layers in enumerate(layers_per_gpu):
        for _ in range(n_layers):
            device_map[f"model.layers.{layer_idx}"] = gpu_id
            layer_idx += 1
    device_map["model.norm"] = 7
    device_map["lm_head"] = 7

    model = dispatch_model(model, device_map=device_map)
    model.eval()

    t1 = time.time()
    logger.info(f"Model loaded, time: {t1 - t0:.2f}s")

    with torch.no_grad():
        logger.info("Running wikitext PPL evaluation...")
        t0 = time.time()
        calculate_ppl_for_wikitext(model, tokenizer, all_input_ids)
        t1 = time.time()
        logger.info(f"Wikitext PPL evaluation total time: {t1 - t0:.2f}s")