diff --git a/benchmark_all.sh b/benchmark_all.sh new file mode 100755 index 0000000..90440c1 --- /dev/null +++ b/benchmark_all.sh @@ -0,0 +1,132 @@ +#!/bin/bash +# Benchmarking script for Fast-dLLM models using the PLENA eval_dllm harness. +# Streams output to the terminal; no per-run log files are retained. +# Only the tabulated results (all_results.jsonl) persist. + +export TRANSFORMERS_OFFLINE=1 +export PYTHONPATH=$PYTHONPATH:/home/gm1425/Fast-dLLM/PLENA_Software/.uv_cache/git-v0/checkouts/c53f323a805b777e/cb963b8/src/ + +LOG_DIR="logs/benchmarks" +mkdir -p $LOG_DIR + +MODELS=( + "Efficient-Large-Model/Fast_dLLM_v2_7B" +) + +# Sweep dimensions are env-overridable (space-separated) so the sweep can be +# sharded across GPUs, e.g.: +# GPU=1 BENCH_BL="8" BENCH_BS="64 128" RESULTS_FILE=logs/benchmarks/gpu1.jsonl bash benchmark_all.sh +# Larger batches (64/128) saturate compute on the eager-quantized configs; +# memory is dominated by the ~15GB of weights and barely grows with batch. +BLOCK_LENGTHS=(${BENCH_BL:-8 16 32}) +BATCH_SIZES=(${BENCH_BS:-16 32}) + +# Active sweep: configs 1-3 run with no external prerequisites. +# 01 fp16 — baseline (no quant_config) +# 02 w4 RTN — weight-only INT4 (Linear quant) +# 03 w4+act4+kv4 — adds self_attn module quant (qk/av/kv-cache INT4). +# Needs the Fast_dLLM_QwenAttentionMXInt support added to +# chop (chop/nn/quantized/modules/fast_dllm_qwen). +# +# Configs 04-06 are GPTQ and are DEFERRED — they require, in addition to the +# chop support above: +# - calib/Qwen_Qwen3-8B_gsm8k_n64_s1024.pt (Step 0 calibrate; see +# plena_experiments/table9/README.md), which needs Qwen/Qwen3-8B (online). +# - a writable gptq.checkpoint_dir (the TOMLs point at /data/models/cx922, +# owned by another user — override to a local path before enabling). +# - config 06 additionally reuses config 05's rotation-search checkpoint. +# Move them into CONFIGS once those prerequisites are in place. +CONFIGS=( + # "plena_experiments/table9/configs/gsm8k/01_fp16.toml" + # "plena_experiments/table9/configs/gsm8k/02_w4_rtn.toml" + # "plena_experiments/table9/configs/gsm8k/03_w4_act4_kv4_rtn.toml" + # "plena_experiments/table9/configs/gsm8k/04_w4_act4_kv4_gptq.toml" + # "plena_experiments/table9/configs/gsm8k/05_w4_act4_kv4_gptq_erryclip.toml" + "plena_experiments/table9/configs/gsm8k/06_w4_act4_kv4_gptq_erryclip_selrot.toml" +) +# Override the config set via BENCH_CONFIGS (space-separated TOML paths). +if [ -n "${BENCH_CONFIGS:-}" ]; then + CONFIGS=(${BENCH_CONFIGS}) +fi + +RESULTS_FILE="${RESULTS_FILE:-$LOG_DIR/all_results.jsonl}" +if [ ! -f "$RESULTS_FILE" ]; then + echo "[]" > "$RESULTS_FILE" +fi + +DEFAULT_PY=$([[ -x .venv/bin/python ]] && echo .venv/bin/python || echo python) +PY=${PY:-$DEFAULT_PY} + +for MODEL in "${MODELS[@]}"; do + for CONFIG in "${CONFIGS[@]}"; do + for BL in "${BLOCK_LENGTHS[@]}"; do + for BS in "${BATCH_SIZES[@]}"; do + MODEL_NAME=$(basename "$MODEL") + CONFIG_NAME=$(basename "$CONFIG" .toml) + RUN_ID="${MODEL_NAME}_${CONFIG_NAME}_BL${BL}_BS${BS}" + + if grep -q "\"run_id\": \"$RUN_ID\"" "$RESULTS_FILE"; then + echo "Skipping $RUN_ID (already exists)..." + continue + fi + + echo "Benchmarking: $MODEL_NAME config=$CONFIG_NAME BL=$BL BS=$BS..." + + # Build args for PLENA eval_dllm jsonargparse CLI. + # No --log_dir: skip persistent per-run args.json/results.json. + ARGS=( + --model_name "$MODEL" + --tasks gsm8k + --device_id cuda:0 + --batch_size "$BS" + --bd_size "$BL" + --small_block_size 8 + --max_new_tokens 512 + --threshold 1.0 + --temperature 0.0 + --num_fewshot 0 + --show_speed true + ) + + # Baseline fp16 runs: no quant_config; quantized runs: pass config + if [[ "$CONFIG_NAME" != *fp16* ]]; then + ARGS+=(--quant_config "$CONFIG") + fi + + # Stream to terminal while capturing to a temp file for metric + # extraction; the temp file is deleted once results are tabulated. + TMP_LOG=$(mktemp) + CUDA_VISIBLE_DEVICES=2 $PY -m quant_eval.cli.eval_dllm "${ARGS[@]}" \ + 2>&1 | tee "$TMP_LOG" + + # gsm8k reports two exact_match scores: "strict-match" (needs a + # literal "#### ", which our \boxed{} prompt never produces -> + # always ~0) and "flexible-extract" (last number in the output -> + # the real score). Prefer flexible-extract; fall back to the last + # exact_match line for tasks (e.g. minerva_math) with no filters. + ACCURACY=$(grep -oP 'exact_match,flexible-extract: \K[0-9.]+' "$TMP_LOG" | head -n 1) + [ -z "$ACCURACY" ] && ACCURACY=$(grep -oP 'exact_match[^,]*?: \K[0-9.]+' "$TMP_LOG" | tail -n 1) + TOKENS_PER_SEC=$(grep -oP 'Tokens/s: \K[0-9.]+' "$TMP_LOG" | tail -n 1) + TTFT_MS=$(grep -oP 'TTFT: \K[0-9.]+' "$TMP_LOG" | tail -n 1) + TOTAL_TIME=$(grep -oP 'Total workload time: \K[0-9.]+' "$TMP_LOG" | tail -n 1) + + rm -f "$TMP_LOG" + + # Treat a run as failed (crash/OOM) if no speed line was emitted. + # Don't record failed runs, so the skip-check retries them later. + if [ -z "$TOKENS_PER_SEC" ]; then + echo "FAILED $RUN_ID — no result (crash/incomplete); will retry on rerun" + continue + fi + + [ -z "$ACCURACY" ] && ACCURACY=0 + [ -z "$TTFT_MS" ] && TTFT_MS=0 + [ -z "$TOTAL_TIME" ] && TOTAL_TIME=0 + + echo "{\"run_id\": \"$RUN_ID\", \"model\": \"$MODEL_NAME\", \"config\": \"$CONFIG_NAME\", \"block_length\": $BL, \"batch_size\": $BS, \"accuracy\": $ACCURACY, \"tokens_per_sec\": $TOKENS_PER_SEC, \"ttft_ms\": $TTFT_MS, \"total_time\": $TOTAL_TIME}" >> "$RESULTS_FILE" + + echo "Done $RUN_ID — Acc: $ACCURACY Tps: $TOKENS_PER_SEC TTFT: ${TTFT_MS}ms" + done + done + done +done diff --git a/benchmark_memory.sh b/benchmark_memory.sh new file mode 100755 index 0000000..578d6fc --- /dev/null +++ b/benchmark_memory.sh @@ -0,0 +1,158 @@ +#!/bin/bash +# Memory profiling for Fast-dLLM models using the PLENA profile_memory CLI. +# +# For each quant config the model is loaded ONCE, then a single generation +# batch is run for every (block_length, batch_size) cell to record: +# * model_footprint_mb — resident model weights (CUDA allocated after load) +# * peak_memory_mb — peak CUDA allocation while generating one batch +# +# Results are merged into the SAME all_results.jsonl produced by +# benchmark_all.sh, matched by run_id (latency/accuracy fields are preserved; +# the memory fields are added). Cells with no existing run_id are appended. + +export TRANSFORMERS_OFFLINE=1 +export PYTHONPATH=$PYTHONPATH:/home/gm1425/Fast-dLLM/PLENA_Software/.uv_cache/git-v0/checkouts/c53f323a805b777e/cb963b8/src/ + +LOG_DIR="logs/benchmarks" +mkdir -p $LOG_DIR + +MODELS=( + "Efficient-Large-Model/Fast_dLLM_v2_7B" +) + +BLOCK_LENGTHS=(8 16 32 64) +BATCH_SIZES=(1 4 8 16 32) + +CONFIGS=( + "plena_experiments/table9/configs/gsm8k/01_fp16.toml" + "plena_experiments/table9/configs/gsm8k/02_w4_rtn.toml" + "plena_experiments/table9/configs/gsm8k/03_w4_act4_kv4_rtn.toml" + "plena_experiments/table9/configs/gsm8k/04_w4_act4_kv4_gptq.toml" + "plena_experiments/table9/configs/gsm8k/05_w4_act4_kv4_gptq_erryclip.toml" + "plena_experiments/table9/configs/gsm8k/06_w4_act4_kv4_gptq_erryclip_selrot.toml" +) + +RESULTS_FILE="$LOG_DIR/all_results.jsonl" +if [ ! -f "$RESULTS_FILE" ]; then + echo "[]" > "$RESULTS_FILE" +fi + +DEFAULT_PY=$([[ -x .venv/bin/python ]] && echo .venv/bin/python || echo python) +PY=${PY:-$DEFAULT_PY} + +# JSON list literals for the profiler CLI (jsonargparse). +BL_JSON=$(printf '%s,' "${BLOCK_LENGTHS[@]}"); BL_JSON="[${BL_JSON%,}]" +BS_JSON=$(printf '%s,' "${BATCH_SIZES[@]}"); BS_JSON="[${BS_JSON%,}]" + +for MODEL in "${MODELS[@]}"; do + for CONFIG in "${CONFIGS[@]}"; do + MODEL_NAME=$(basename "$MODEL") + CONFIG_NAME=$(basename "$CONFIG" .toml) + + echo "Profiling memory: $MODEL_NAME config=$CONFIG_NAME (single batch per cell)..." + + ARGS=( + --model_name "$MODEL" + --device_id cuda:0 + --bd_sizes "$BL_JSON" + --batch_sizes "$BS_JSON" + --small_block_size 8 + --max_new_tokens 512 + --threshold 1.0 + --temperature 0.0 + ) + if [[ "$CONFIG_NAME" != *fp16* ]]; then + ARGS+=(--quant_config "$CONFIG") + fi + + TMP_LOG=$(mktemp) + CUDA_VISIBLE_DEVICES=3 $PY -m quant_eval.cli.profile_memory "${ARGS[@]}" \ + 2>&1 | tee "$TMP_LOG" + + if ! grep -q "MEM_RESULT" "$TMP_LOG"; then + echo "FAILED $MODEL_NAME/$CONFIG_NAME — no MEM_RESULT lines (crash/incomplete)" + rm -f "$TMP_LOG" + continue + fi + + # Merge each MEM_RESULT cell into all_results.jsonl by run_id. + "$PY" - "$RESULTS_FILE" "$TMP_LOG" "$MODEL_NAME" "$CONFIG_NAME" <<'PYEOF' +import json, re, sys + +results_path, log_path, model_name, config_name = sys.argv[1:5] + +# Parse memory cells from the profiler log. +cells = {} # run_id -> (footprint_mb, peak_mb, peak_reserved_mb, bl, bs) +pat = re.compile( + r"MEM_RESULT\|bd=(\d+)\|bs=(\d+)\|cache=(true|false)\|footprint_mb=([\d.]+)" + r"\|peak_mb=([\d.]+)\|peak_reserved_mb=([\d.]+)" +) +with open(log_path) as f: + for line in f: + m = pat.search(line) + if not m: + continue + bl, bs, cache = int(m.group(1)), int(m.group(2)), m.group(3) + # run_id must match benchmark_all.sh exactly: ..._BL{bl}_BS{bs}_CACHE{cache} + run_id = f"{model_name}_{config_name}_BL{bl}_BS{bs}_CACHE{cache}" + cells[run_id] = { + "block_length": bl, + "batch_size": bs, + "use_block_cache": cache == "true", + "model_footprint_mb": round(float(m.group(4)), 2), + "peak_memory_mb": round(float(m.group(5)), 2), + "peak_reserved_mb": round(float(m.group(6)), 2), + } + +# Read existing lines, updating matching records in place. +out_lines = [] +seen = set() +with open(results_path) as f: + for raw in f: + s = raw.rstrip("\n") + if not s.strip(): + continue + try: + obj = json.loads(s) + except json.JSONDecodeError: + out_lines.append(s) + continue + if isinstance(obj, dict) and obj.get("run_id") in cells: + rid = obj["run_id"] + obj.update( + model_footprint_mb=cells[rid]["model_footprint_mb"], + peak_memory_mb=cells[rid]["peak_memory_mb"], + peak_reserved_mb=cells[rid]["peak_reserved_mb"], + ) + seen.add(rid) + out_lines.append(json.dumps(obj)) + else: + out_lines.append(s) + +# Append cells that had no existing latency/accuracy record. +for rid, c in cells.items(): + if rid in seen: + continue + out_lines.append(json.dumps({ + "run_id": rid, + "model": model_name, + "config": config_name, + "block_length": c["block_length"], + "batch_size": c["batch_size"], + "use_block_cache": c["use_block_cache"], + "model_footprint_mb": c["model_footprint_mb"], + "peak_memory_mb": c["peak_memory_mb"], + "peak_reserved_mb": c["peak_reserved_mb"], + })) + +with open(results_path, "w") as f: + f.write("\n".join(out_lines) + "\n") + +print(f"Merged {len(cells)} memory cells " + f"({len(seen)} updated, {len(cells) - len(seen)} appended) for {config_name}.") +PYEOF + + rm -f "$TMP_LOG" + echo "Done memory profiling $MODEL_NAME/$CONFIG_NAME." + done +done diff --git a/plena_experiments/table9/configs/gsm8k/04_w4_act4_kv4_gptq.toml b/plena_experiments/table9/configs/gsm8k/04_w4_act4_kv4_gptq.toml index 96c6f01..a809fde 100644 --- a/plena_experiments/table9/configs/gsm8k/04_w4_act4_kv4_gptq.toml +++ b/plena_experiments/table9/configs/gsm8k/04_w4_act4_kv4_gptq.toml @@ -6,15 +6,15 @@ by = "regex_name" [gptq] -model_name = "Qwen/Qwen3-8B" -dataset = "file:calib/Qwen_Qwen3-8B_gsm8k_n64_s1024.pt" +model_name = "Efficient-Large-Model/Fast_dLLM_v2_7B" +dataset = "file:calib/Fast_dLLM_v2_7B_gsm8k_n64_s1024.pt" nsamples = 64 seqlen = 1024 format = "mxint" quantile_search = true clip_search_y = false cali_batch_size = 32 -checkpoint_dir = "/data/models/cx922/plena_camera/gsm8k_04_w4_act4_kv4_gptq" +checkpoint_dir = "/data/models/gm1425/plena_gptq/fastdllm_gsm8k_04_w4_act4_kv4_gptq" [gptq.weight_config] weight_block_size = 16 diff --git a/plena_experiments/table9/configs/gsm8k/05_w4_act4_kv4_gptq_erryclip.toml b/plena_experiments/table9/configs/gsm8k/05_w4_act4_kv4_gptq_erryclip.toml index 652e9e6..3b64465 100644 --- a/plena_experiments/table9/configs/gsm8k/05_w4_act4_kv4_gptq_erryclip.toml +++ b/plena_experiments/table9/configs/gsm8k/05_w4_act4_kv4_gptq_erryclip.toml @@ -6,15 +6,15 @@ by = "regex_name" [gptq] -model_name = "Qwen/Qwen3-8B" -dataset = "file:calib/Qwen_Qwen3-8B_gsm8k_n64_s1024.pt" +model_name = "Efficient-Large-Model/Fast_dLLM_v2_7B" +dataset = "file:calib/Fast_dLLM_v2_7B_gsm8k_n64_s1024.pt" nsamples = 64 seqlen = 1024 format = "mxint" quantile_search = true clip_search_y = true # ← the only diff from 04 cali_batch_size = 32 -checkpoint_dir = "/data/models/cx922/plena_camera/gsm8k_05_w4_act4_kv4_gptq_erryclip" +checkpoint_dir = "/data/models/gm1425/plena_gptq/fastdllm_gsm8k_05_w4_act4_kv4_gptq_erryclip" [gptq.weight_config] weight_block_size = 16 diff --git a/plena_experiments/table9/configs/gsm8k/06_w4_act4_kv4_gptq_erryclip_selrot.toml b/plena_experiments/table9/configs/gsm8k/06_w4_act4_kv4_gptq_erryclip_selrot.toml index 8eabd66..327ed87 100644 --- a/plena_experiments/table9/configs/gsm8k/06_w4_act4_kv4_gptq_erryclip_selrot.toml +++ b/plena_experiments/table9/configs/gsm8k/06_w4_act4_kv4_gptq_erryclip_selrot.toml @@ -11,15 +11,15 @@ by = "regex_name" [gptq] -model_name = "Qwen/Qwen3-8B" -dataset = "file:calib/Qwen_Qwen3-8B_gsm8k_n64_s1024.pt" +model_name = "Efficient-Large-Model/Fast_dLLM_v2_7B" +dataset = "file:calib/Fast_dLLM_v2_7B_gsm8k_n64_s1024.pt" nsamples = 64 seqlen = 1024 format = "mxint" quantile_search = true clip_search_y = true cali_batch_size = 32 -checkpoint_dir = "/data/models/cx922/plena_camera/gsm8k_05_w4_act4_kv4_gptq_erryclip" +checkpoint_dir = "/data/models/gm1425/plena_gptq/fastdllm_gsm8k_05_w4_act4_kv4_gptq_erryclip" [gptq.weight_config] weight_block_size = 16 diff --git a/quant_eval/cli/eval_dllm.py b/quant_eval/cli/eval_dllm.py index 53e2ad6..8f4c8cd 100644 --- a/quant_eval/cli/eval_dllm.py +++ b/quant_eval/cli/eval_dllm.py @@ -36,17 +36,51 @@ create_experiment_log_dir, save_args, save_results, + _fix_rotary_buffers, ) from quant_eval.eval.dllm_v2.dllm_generation import setup_dllm_generation from quant_eval.eval.dllm_v2.eval_dllm import evaluate_dllm from quant_eval.quantize import load_quant_config logger = get_logger(__name__) -set_logging_verbosity("debug") +set_logging_verbosity("info") + + +def _normalize_tied_weights_keys(model: torch.nn.Module) -> None: + """Coerce any list-valued ``_tied_weights_keys`` to the dict form that + newer transformers expects. + + Older models (e.g. Fast-dLLM's remote ``modeling.py``) declare + ``_tied_weights_keys = ["lm_head.weight"]``. The new ``save_pretrained`` + path iterates ``_tied_weights_keys.keys()``, so a list raises + ``AttributeError: 'list' object has no attribute 'keys'``. The dict maps + each tied (alias) weight to the source weight it shares storage with — + here the input-embedding weight. + """ + for submodule in model.modules(): + tied = getattr(submodule, "_tied_weights_keys", None) + if not isinstance(tied, (list, tuple)): + continue + source = None + get_input_embeddings = getattr(submodule, "get_input_embeddings", None) + if callable(get_input_embeddings): + try: + emb = get_input_embeddings() + except (AttributeError, NotImplementedError): + emb = None + if emb is not None: + # Find the dotted name of the embedding's weight within submodule. + for name, param in submodule.named_parameters(): + if param is getattr(emb, "weight", None): + source = name + break + # Fall back to self-mapping if we can't resolve the source; this still + # satisfies the `.keys()` access and tying is re-applied from config. + submodule._tied_weights_keys = {k: (source or k) for k in tied} def main( - model_name: str = "Efficient-Large-Model/Fast_dLLM_v2_1.5B", + model_name: str = "Efficient-Large-Model/Fast_dLLM_v2_7B", tasks: Union[str, list[str]] = "gsm8k", device_id: str = "cuda:0", dtype: str = "bfloat16", @@ -59,7 +93,10 @@ def main( mask_id: int = 151665, bd_size: int = 32, small_block_size: int = 8, - threshold: float = 1.0, + threshold: float = 0.9, + use_block_cache: bool = False, + temperature: float = 0.0, + top_p: float = 0.95, show_speed: bool = True, log_dir: Union[str, None] = None, ): @@ -89,6 +126,9 @@ def main( small_block_size: Inner block size for iterative unmasking within each outer block. threshold: Confidence threshold for committing unmasked tokens. + use_block_cache: Cache intermediate block KV states for faster decoding. + temperature: Sampling temperature. + top_p: Top-p sampling probability. show_speed: Log throughput metrics (tokens/second). log_dir: Directory for ``args.json`` and ``results.json``. @@ -101,6 +141,7 @@ def main( print(f"Model: {model_name}") print(f"Tasks: {tasks}") print(f"Block size: {bd_size}, Sub-block: {small_block_size}, Threshold: {threshold}") + print(f"Block cache: {use_block_cache}, Temperature: {temperature}, Top-p: {top_p}") quantize = quant_config is not None if quantize: @@ -109,6 +150,7 @@ def main( print("Quantization: None (baseline)") print("=" * 60) + from pathlib import Path if log_dir: log_dir = create_experiment_log_dir(log_dir) save_args(log_dir, locals().copy()) @@ -133,21 +175,84 @@ def main( ) model.eval() - if quantize: - from chop.passes.module.transforms import quantize_module_transform_pass + # Place the model on GPU *before* quantization so the loaded weights and the + # quant passes (GPTQ / rotation-search calibration) all run on-device rather + # than on CPU — setup_model loads onto CPU. For model-parallel we defer to + # the dispatch_model call after quantization, since accelerate's offload + # hooks don't compose with the in-place module-replacement passes. + if not model_parallel: + logger.info("Moving model to %s", device_id) + model.to(device_id) + if quantize: pass_args = load_quant_config(quant_config) - if "gptq" in pass_args: - pass_args["gptq"]["device"] = device_id - - n_linear = sum( - 1 for _, m in model.named_modules() if isinstance(m, torch.nn.Linear) + + # Check if we're running Table 9 configuration 06. Config 06 is the most + # expensive because it adds rotation search on top of GPTQ + erryclip, + # so we optionally cache it to a separate filesystem (/tmp). + is_config6 = "table9" in str(quant_config) and "06_" in Path(quant_config).name + cache_dir = Path("/tmp/plena_quant_cache") / f"{model_name.replace('/', '--')}_table9_config6" if is_config6 else None + + loaded_from_cache = False + if cache_dir and cache_dir.exists(): + logger.info("Reloading pre-quantized model from %s", cache_dir) + model = transformers.AutoModelForCausalLM.from_pretrained( + cache_dir, trust_remote_code=True, torch_dtype=torch_dtype + ) + _fix_rotary_buffers(model, logger) + if not model_parallel: + model.to(device_id) + loaded_from_cache = True + + if not loaded_from_cache: + from chop.passes.module.transforms import quantize_module_transform_pass + + # Plumb device (+ model_name) into the GPTQ / rotation-search passes the + # same way eval_lm does — the MASE passes need these but they don't + # belong in the TOML schema. Required for configs 04-06 (GPTQ) and 06 + # (rotation_search), so all six table9 rows run through this CLI. + if "gptq" in pass_args: + pass_args["gptq"]["device"] = device_id + if "rotation_search" in pass_args: + pass_args["rotation_search"]["device"] = device_id + pass_args["rotation_search"].setdefault("model_name", model_name) + + n_linear = sum( + 1 for _, m in model.named_modules() if isinstance(m, torch.nn.Linear) + ) + logger.info("Quantizing %d linear layers...", n_linear) + t0 = time.time() + model, _ = quantize_module_transform_pass(model, pass_args) + logger.info("Quantization complete in %.1fs", time.time() - t0) + + if cache_dir: + logger.info("Saving quantized model to %s", cache_dir) + cache_dir.mkdir(parents=True, exist_ok=True) + # Fast-dLLM's remote modeling.py declares `_tied_weights_keys` + # as a list (old transformers convention), but the installed + # transformers' save path calls `.keys()` on it (new dict + # convention). Normalize any list-valued attr to a dict mapping + # each tied weight -> the input-embedding weight it aliases. + _normalize_tied_weights_keys(model) + model.save_pretrained(cache_dir) + + # Surface which classes the dispatch landed on, so you can confirm the + # rotate / GPTQ variants are wired in when the TOML asks for them. + from collections import Counter + + cls_count = Counter( + type(m).__name__ for _, m in model.named_modules() + if "MX" in type(m).__name__ ) - logger.info("Quantizing %d linear layers...", n_linear) - t0 = time.time() - model, _ = quantize_module_transform_pass(model, pass_args) - logger.info("Quantization complete in %.1fs", time.time() - t0) - + if cls_count: + logger.info( + "Post-quant module classes:\n%s", + "\n".join(f" {c}: {n}" for c, n in cls_count.most_common()), + ) + + # Final device placement. For model-parallel this is the real dispatch step; + # for single-GPU it's a safety net that catches any modules the quant passes + # created on CPU (the model is already on device_id from before quantization). if model_parallel: model = move_to_gpu(model, model_parallel) else: @@ -173,9 +278,12 @@ def main( bd_size=bd_size, small_block_size=small_block_size, threshold=threshold, + use_block_cache=use_block_cache, + temperature=temperature, show_speed=show_speed, ) + print("\n" + "=" * 60) print("Results:") print("=" * 60) diff --git a/quant_eval/cli/eval_evalplus.py b/quant_eval/cli/eval_evalplus.py index 1418759..febe2f6 100644 --- a/quant_eval/cli/eval_evalplus.py +++ b/quant_eval/cli/eval_evalplus.py @@ -43,11 +43,11 @@ def main( - model_name: str = "Qwen/Qwen2.5-1.5B", + model_name: str = "Efficient-Large-Model/Fast_dLLM_v2_7B", dataset: str = "humaneval", device_id: str = "cuda:0", dtype: str = "bfloat16", - quant_config: Union[str, None] = "quant_eval/configs/llama_mxint4.toml", + quant_config: Union[str, None] = None, model_parallel: bool = False, batch_size: int = 1, greedy: bool = False, @@ -59,6 +59,11 @@ def main( parallel: Union[int, None] = None, version: str = "default", log_dir: Union[str, None] = None, + # dLLM specific + mask_id: int = 151665, + bd_size: int = 32, + small_block_size: int = 8, + threshold: float = 0.9, ): """ Run evalplus (HumanEval+ / MBPP+) on an optionally MX-quantized HF model. @@ -83,6 +88,10 @@ def main( runs serially. version: evalplus dataset version (e.g. ``"default"``). log_dir: Directory for ``args.json`` and ``results.json``. + mask_id: Token ID used as the diffusion mask (for dLLM). + bd_size: Outer block-diffusion block size (for dLLM). + small_block_size: Inner block size for iterative unmasking (for dLLM). + threshold: Confidence threshold for committing tokens (for dLLM). Returns: evalplus results dict — pass@1 (and pass@k when applicable) plus diff --git a/quant_eval/cli/eval_llada.py b/quant_eval/cli/eval_llada.py index 9381e05..5150bb9 100644 --- a/quant_eval/cli/eval_llada.py +++ b/quant_eval/cli/eval_llada.py @@ -1,22 +1,15 @@ """ -LLaDA (diffusion-style language model) evaluation with optional MX quantization. +Fast-dLLM v2 evaluation via the legacy ``llada_dist`` lm-eval alias. -Wraps lm-eval-harness's CLI. Use ``--model llada_dist`` and pass model and -quantization options through lm-eval's ``--model_args`` flag. +Use ``--model llada_dist`` and pass Fast-dLLM v2 options through lm-eval's +``--model_args`` flag. -Example — baseline (prefix cache): +Example: python -m quant_eval.cli.eval_llada \\ --tasks gsm8k --num_fewshot 0 \\ --model llada_dist \\ - --model_args model_path='GSAI-ML/LLaDA-8B-Instruct',gen_length=256,steps=256,block_length=32,use_cache=True - -Example — with MXINT4 KV-cache quantization: - - python -m quant_eval.cli.eval_llada \\ - --tasks gsm8k --num_fewshot 0 \\ - --model llada_dist \\ - --model_args model_path='GSAI-ML/LLaDA-8B-Instruct',gen_length=256,steps=256,block_length=32,use_cache=True,quant_config='quant_eval/configs/llama_mxint4.toml' + --model_args model_path='Efficient-Large-Model/Fast_dLLM_v2_7B',gen_length=256,steps=256,block_length=32,use_cache=True """ # Import to trigger @register_model("llada_dist") diff --git a/quant_eval/cli/profile_memory.py b/quant_eval/cli/profile_memory.py new file mode 100644 index 0000000..c0226ce --- /dev/null +++ b/quant_eval/cli/profile_memory.py @@ -0,0 +1,196 @@ +""" +Memory footprint / peak-usage profiler for Fast-dLLM v2. + +Loads a model **once** (optionally quantized) and sweeps a grid of block sizes +and batch sizes, running a single generation batch per cell and recording: + + * ``model_footprint_mb`` — CUDA memory allocated by the resident model weights + (measured once after load; identical across cells for a given quant config). + * ``peak_memory_mb`` — peak CUDA memory allocated while generating one + batch (weights + activations + KV cache), isolated per cell. + +Each cell is emitted on stdout as a single parseable line, e.g.:: + + MEM_RESULT|bd=16|bs=4|cache=false|footprint_mb=14820.50|peak_mb=15960.20|peak_reserved_mb=16210.00 + +``benchmark_memory.sh`` consumes these lines and merges them into +``all_results.jsonl`` by ``run_id``. + +Example:: + + python -m quant_eval.cli.profile_memory \\ + --model_name Efficient-Large-Model/Fast_dLLM_v2_7B \\ + --bd_sizes '[8,16,32,64]' --batch_sizes '[1,4,8,16,32]' +""" + +from typing import List, Union +import time + +import torch +import transformers +from datasets import load_dataset + +from quant_eval.utils import get_logger, set_logging_verbosity, setup_model +from quant_eval.eval.dllm_v2.dllm_generation import ( + setup_dllm_generation, + FAST_DLLM_MASK_ID, +) +from quant_eval.quantize import load_quant_config + +logger = get_logger(__name__) +set_logging_verbosity("info") + +MB = 1024.0 * 1024.0 + + +def _build_batch(tokenizer, questions, mask_id, device): + """Tokenize gsm8k-style prompts and right-pad with the mask id, mirroring + the eval harness's ``generate_until`` batching.""" + encoded = [] + for q in questions: + text = ("Question: " + q + "\nAnswer:").replace( + "Answer:", + "Please reason step by step, and put your final answer within \\boxed{}.", + ) + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": text}], + add_generation_prompt=True, + tokenize=False, + ) + encoded.append(tokenizer([prompt], return_tensors="pt").to(device)["input_ids"]) + + max_len = max(e.shape[1] for e in encoded) + seq_len = [e.shape[1] for e in encoded] + padded = [ + torch.cat( + [ + e, + torch.full( + (1, max_len - e.shape[1]), mask_id, dtype=torch.long, device=device + ), + ], + dim=1, + ) + for e in encoded + ] + batch_ids = torch.cat(padded, dim=0) + return batch_ids, seq_len + + +def main( + model_name: str = "Efficient-Large-Model/Fast_dLLM_v2_7B", + quant_config: Union[str, None] = None, + device_id: str = "cuda:0", + dtype: str = "bfloat16", + bd_sizes: List[int] = [8, 16, 32, 64], + batch_sizes: List[int] = [1, 4, 8, 16, 32], + small_block_size: int = 8, + max_new_tokens: int = 512, + threshold: float = 1.0, + temperature: float = 0.0, + block_caches: List[bool] = [False, True], + mask_id: int = FAST_DLLM_MASK_ID, +): + """ + Profile model footprint and peak generation memory across a bd_size × + batch_size grid. The model is loaded once; each grid cell runs a single + generation batch with peak-memory stats reset beforehand so cells don't + contaminate each other. + """ + print("=" * 60) + print("Fast-dLLM Memory Profiler") + print(f"Model: {model_name}") + print(f"Quant: {quant_config or 'None (baseline)'}") + print(f"bd_sizes={bd_sizes} batch_sizes={batch_sizes}") + print("=" * 60) + + transformers.set_seed(0) + dtype_map = { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, + } + torch_dtype = dtype_map.get(dtype, torch.bfloat16) + device = torch.device(device_id) + # Initialize the CUDA context before touching memory stats — the memory + # APIs raise "Invalid device argument" if called before any CUDA init. + torch.cuda.set_device(device) + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats(device) + + tokenizer, model = setup_model( + model_name, False, dtype=torch_dtype, device=device_id + ) + model.eval() + + if quant_config is not None: + from chop.passes.module.transforms import quantize_module_transform_pass + + pass_args = load_quant_config(quant_config) + if "gptq" in pass_args: + pass_args["gptq"]["device"] = device_id + t0 = time.time() + model, _ = quantize_module_transform_pass(model, pass_args) + logger.info("Quantization complete in %.1fs", time.time() - t0) + + model.to(device_id) + setup_dllm_generation(model) + + # Resident model footprint: allocated memory once weights are settled. + torch.cuda.synchronize(device) + footprint_mb = torch.cuda.memory_allocated(device) / MB + param_mb = sum(p.numel() * p.element_size() for p in model.parameters()) / MB + buffer_mb = sum(b.numel() * b.element_size() for b in model.buffers()) / MB + print( + f"Model footprint: {footprint_mb:.2f} MB allocated " + f"(params {param_mb:.2f} MB + buffers {buffer_mb:.2f} MB)" + ) + + # A pool of real gsm8k prompts to draw batches from (largest bs needed). + ds = load_dataset("openai/gsm8k", "main", split="test") + pool = [ds[i]["question"] for i in range(max(batch_sizes))] + + for bd in bd_sizes: + for bs in batch_sizes: + for bc in block_caches: + batch_ids, seq_len = _build_batch(tokenizer, pool[:bs], mask_id, device) + min_len = min(seq_len) + + torch.cuda.empty_cache() + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + + try: + with torch.no_grad(): + model.mdm_sample( + batch_ids, + tokenizer=tokenizer, + block_size=bd, + small_block_size=small_block_size, + max_new_tokens=max_new_tokens, + mask_id=mask_id, + min_len=int(min_len), + seq_len=torch.tensor(seq_len, device=device), + use_block_cache=bc, + threshold=threshold, + temperature=temperature, + ) + torch.cuda.synchronize(device) + peak_mb = torch.cuda.max_memory_allocated(device) / MB + peak_reserved_mb = torch.cuda.max_memory_reserved(device) / MB + print( + f"MEM_RESULT|bd={bd}|bs={bs}|cache={str(bc).lower()}" + f"|footprint_mb={footprint_mb:.2f}" + f"|peak_mb={peak_mb:.2f}|peak_reserved_mb={peak_reserved_mb:.2f}" + ) + except torch.cuda.OutOfMemoryError: + torch.cuda.empty_cache() + print(f"MEM_OOM|bd={bd}|bs={bs}|cache={str(bc).lower()} — skipped (out of memory)") + + print("\n[INFO] Memory profiling complete.") + + +if __name__ == "__main__": + from jsonargparse import CLI + + CLI(main) \ No newline at end of file diff --git a/quant_eval/eval/dllm_v2/dllm_generation.py b/quant_eval/eval/dllm_v2/dllm_generation.py index cc42d10..ac1b163 100644 --- a/quant_eval/eval/dllm_v2/dllm_generation.py +++ b/quant_eval/eval/dllm_v2/dllm_generation.py @@ -5,6 +5,8 @@ Based on: https://github.com/ML-GSAI/SMDM """ +import time + import torch import types from transformers import DynamicCache @@ -35,6 +37,12 @@ def batch_sample( num_blocks = max_new_tokens // block_size + seq_len.max().item() // block_size batch_size = input_ids.shape[0] + # Time-to-first-token: wall-clock from generation start to the first + # committed (unmasked) token. Reset per call; read by the harness. + torch.cuda.synchronize() if input_ids.is_cuda else None + _ttft_start = time.time() + self.mdm_ttft = None + if min_len > block_size: output = self.forward( input_ids=input_ids[:, :(min_len // block_size * block_size)], @@ -144,6 +152,10 @@ def batch_sample( x_t[:, start:end][unmask_idx] = x_1[unmask_idx] + if self.mdm_ttft is None and unmask_idx.any(): + torch.cuda.synchronize() if x_t.is_cuda else None + self.mdm_ttft = time.time() - _ttft_start + finished_row_flags = ((x_1 == stop_token) & unmask_idx).any(dim=1) finished_flag = finished_flag | finished_row_flags diff --git a/quant_eval/eval/dllm_v2/eval_dllm.py b/quant_eval/eval/dllm_v2/eval_dllm.py index d69992c..12770e3 100644 --- a/quant_eval/eval/dllm_v2/eval_dllm.py +++ b/quant_eval/eval/dllm_v2/eval_dllm.py @@ -40,6 +40,8 @@ def __init__( small_block_size: int = 8, bd_size: int = 32, threshold: float = 0.9, + temperature: float = 0.0, + top_p: float = 0.95, ): super().__init__() @@ -55,6 +57,8 @@ def __init__( self.small_block_size = small_block_size self.bd_size = bd_size self.threshold = threshold + self.temperature = temperature + self.top_p = top_p self._rank = 0 self._world_size = 1 @@ -144,6 +148,7 @@ def _tokenize(e): def generate_until(self, requests): output = [None] * len(requests) num_tokens = 0 + ttfts = [] start_time = time.time() requests_with_indices = [(i, req) for i, req in enumerate(requests)] @@ -209,22 +214,31 @@ def generate_until(self, requests): seq_len=torch.tensor(seq_len, device=self.device), use_block_cache=self.use_block_cache, threshold=self.threshold, + temperature=self.temperature, ) + if self.show_speed: + ttft = getattr(self.model, "mdm_ttft", None) + if ttft is not None: + ttfts.append(ttft) + for batch_pos, (orig_idx, req) in enumerate(batch): - generated_answer = self.tokenizer.decode( - generated_ids[batch_pos][seq_len[batch_pos]:], - skip_special_tokens=True, - ) + gen_ids = generated_ids[batch_pos][seq_len[batch_pos]:] + generated_answer = self.tokenizer.decode(gen_ids, skip_special_tokens=True) if self.show_speed: - num_tokens += (generated_ids[batch_pos][seq_len[batch_pos]:] != self.mask_id).sum() + num_tokens += (gen_ids != self.mask_id).sum() output[orig_idx] = generated_answer if self.show_speed: elapsed = time.time() - start_time - print(f"Total tokens: {num_tokens}, Time: {elapsed:.2f}s, Tokens/s: {num_tokens / elapsed:.2f}") + mean_ttft = sum(ttfts) / len(ttfts) if ttfts else float("nan") + print( + f"Total tokens: {num_tokens}, Time: {elapsed:.2f}s, " + f"Tokens/s: {num_tokens / elapsed:.2f}, " + f"TTFT: {mean_ttft * 1000:.1f}ms" + ) return output @@ -242,6 +256,8 @@ def evaluate_dllm( bd_size: int = 32, small_block_size: int = 8, threshold: float = 1.0, + use_block_cache: bool = False, + temperature: float = 0.0, show_speed: bool = True, ) -> Dict: """ @@ -260,6 +276,9 @@ def evaluate_dllm( bd_size: Block diffusion size. small_block_size: Sub-block size. threshold: Unmasking threshold. + use_block_cache: Cache intermediate block KV states for faster decoding. + temperature: Sampling temperature. 0.0 is deterministic greedy + (reference setting); > 0 enables stochastic top-p sampling. show_speed: Show throughput metrics. Returns: @@ -274,10 +293,11 @@ def evaluate_dllm( max_new_tokens=max_new_tokens, batch_size=batch_size, mask_id=mask_id, - use_block_cache=False, + use_block_cache=use_block_cache, small_block_size=small_block_size, bd_size=bd_size, threshold=threshold, + temperature=temperature, ) task_list = [tasks] if isinstance(tasks, str) else tasks diff --git a/quant_eval/eval/evalplus.py b/quant_eval/eval/evalplus.py index a048cb4..a4a1045 100644 --- a/quant_eval/eval/evalplus.py +++ b/quant_eval/eval/evalplus.py @@ -46,6 +46,11 @@ def __init__( tokenizer: PreTrainedTokenizer, dataset: str, force_base_prompt: bool = False, + # dLLM specific + mask_id: int = 151665, + bd_size: int = 32, + small_block_size: int = 8, + threshold: float = 0.9, **kwargs, ): super().__init__(name="preloaded-hf", **kwargs) @@ -55,6 +60,12 @@ def __init__( self.skip_special_tokens = True self.force_base_prompt = force_base_prompt + # dLLM specific + self.mask_id = mask_id + self.bd_size = bd_size + self.small_block_size = small_block_size + self.threshold = threshold + if self.is_direct_completion(): self.eos += extra_eos_for_direct_completion(dataset) else: @@ -91,25 +102,51 @@ def codegen( kwargs["top_k"] = 20 kwargs["temperature"] = self.temperature - stop_sequencer = StopSequencer( - self.model, model_type="causal", tokenizer=self.tokenizer - ) - orig_get_stopping_criteria = self.model._get_stopping_criteria - model = stop_sequencer.register_stop_texts( - stop_texts=self.eos, - input_length=input_tokens.size(-1), - ) + if hasattr(self.model, "mdm_sample"): + # dLLM sampling + batch_size = min(self.batch_size, num_samples) + batched_input_ids = input_tokens.repeat(batch_size, 1) + seq_len = torch.tensor([input_tokens.size(-1)] * batch_size, device=self.device) + + outputs = self.model.mdm_sample( + batched_input_ids, + tokenizer=self.tokenizer, + block_size=self.bd_size, + small_block_size=self.small_block_size, + max_new_tokens=self.max_new_tokens, + mask_id=self.mask_id, + min_len=int(input_tokens.size(-1)), + seq_len=seq_len, + use_block_cache=False, + threshold=self.threshold, + temperature=self.temperature if do_sample else 0.0, + top_p=0.95 if do_sample else 1.0, + ) + # mdm_sample returns a dict or list of generated IDs. + # Convert to tensor if needed. + if isinstance(outputs, dict): + outputs = torch.stack([outputs[i] for i in range(len(outputs))]) + else: + # Standard autoregressive sampling + stop_sequencer = StopSequencer( + self.model, model_type="causal", tokenizer=self.tokenizer + ) + orig_get_stopping_criteria = self.model._get_stopping_criteria + model = stop_sequencer.register_stop_texts( + stop_texts=self.eos, + input_length=input_tokens.size(-1), + ) - outputs = model.generate( - input_tokens, - max_new_tokens=self.max_new_tokens, - do_sample=do_sample, - num_return_sequences=min(self.batch_size, num_samples), - pad_token_id=self.tokenizer.eos_token_id, - **kwargs, - ) - # Restore original method to prevent nested wrapping accumulation. - self.model._get_stopping_criteria = orig_get_stopping_criteria + outputs = model.generate( + input_tokens, + max_new_tokens=self.max_new_tokens, + do_sample=do_sample, + num_return_sequences=min(self.batch_size, num_samples), + pad_token_id=self.tokenizer.eos_token_id, + **kwargs, + ) + # Restore original method to prevent nested wrapping accumulation. + self.model._get_stopping_criteria = orig_get_stopping_criteria gen_strs = self.tokenizer.batch_decode( outputs[:, input_tokens.size(-1) :], @@ -139,6 +176,11 @@ def evaluate_with_evalplus( base_only: bool = False, version: str = "default", overwrite: bool = False, + # dLLM specific + mask_id: int = 151665, + bd_size: int = 32, + small_block_size: int = 8, + threshold: float = 0.9, ) -> Dict: """Generate code and evaluate with evalplus. @@ -155,6 +197,10 @@ def evaluate_with_evalplus( base_only: Only run base tests (skip plus tests). version: Dataset version. overwrite: If True, regenerate even if a previous jsonl exists. + mask_id: Mask token ID for dLLM. + bd_size: Block diffusion size. + small_block_size: Sub-block size. + threshold: Unmasking threshold. Returns: Dictionary with evaluation results. @@ -184,6 +230,10 @@ def evaluate_with_evalplus( max_new_tokens=max_new_tokens, instruction_prefix=instruction_prefix, response_prefix=response_prefix, + mask_id=mask_id, + bd_size=bd_size, + small_block_size=small_block_size, + threshold=threshold, ) if output_dir is None: diff --git a/quant_eval/eval/llada/eval_llada.py b/quant_eval/eval/llada/eval_llada.py index d30d80f..4b3c813 100644 --- a/quant_eval/eval/llada/eval_llada.py +++ b/quant_eval/eval/llada/eval_llada.py @@ -16,7 +16,7 @@ # Modified from LLaDA repos: https://github.com/ML-GSAI/LLaDA # Modified from Fast-dLLM: https://github.com/NVlabs/Fast-dLLM -''' +""" LLaDA evaluation harness with Fast-dLLM v1 KV cache acceleration. Registers a "llada_dist" model with lm-eval. Supports optional MASE @@ -31,9 +31,8 @@ python -m quant_eval.eval.eval_llada --tasks gsm8k --num_fewshot 0 \ --model llada_dist \ --model_args model_path='GSAI-ML/LLaDA-8B-Instruct',gen_length=256,steps=256,block_length=32,use_cache=True,quant_config='configs/kv_only_mxint4.toml' -''' +""" -import accelerate import torch import random import numpy as np @@ -43,18 +42,21 @@ from lm_eval.api.model import LM from lm_eval.api.registry import register_model from tqdm import tqdm -import os import json import time - -from transformers import AutoTokenizer, AutoConfig -from chop.nn.quantized.modules.llada import LLaDAModelLM +import logging from quant_eval.eval.llada.llada_generation import ( generate, generate_with_prefix_cache, generate_with_dual_cache, ) +from quant_eval.quantize import load_quant_config +from quant_eval.utils import get_logger, move_to_gpu, set_logging_verbosity, setup_model + +logger = get_logger(__name__) +logger.setLevel(logging.DEBUG) +set_logging_verbosity("debug") def set_seed(seed): @@ -69,17 +71,15 @@ def set_seed(seed): class LLaDAEvalHarness(LM): def __init__( self, - model_path='', - mask_id=126336, + model_path="", + mask_id=151665, max_length=4096, - batch_size=32, mc_num=128, is_check_greedy=True, steps=1024, gen_length=1024, block_length=1024, - remasking='low_confidence', - device="cuda", + remasking="low_confidence", use_cache=False, threshold=None, factor=None, @@ -90,54 +90,85 @@ def __init__( **kwargs, ): super().__init__() - - accelerator = accelerate.Accelerator() - if accelerator.num_processes > 1: - self.accelerator = accelerator + self.model_path = model_path + + # --------------------------------------------------------------- + # FORCE CORRECT PARAMETERS FOR FAST-DLLM V2 (QWEN-BASED) + # --------------------------------------------------------------- + is_fast_dllm = "fast_dllm" in model_path.lower() or "qwen" in model_path.lower() + + if is_fast_dllm: + self.mask_id = 151665 # Official Fast-dLLM v2 mask token + self.steps = 256 # CRITICAL: Prevents 1024-step bottleneck + self.threshold = 0.8 # Balanced for math reasoning & parallel decoding + self.block_length = max(int(block_length), 16) + logger.info( + f"🔧 Fast-dLLM v2 detected. Forcing: mask_id={self.mask_id}, steps={self.steps}, threshold={self.threshold}, block_length={self.block_length}" + ) else: - self.accelerator = None - - model_kwargs = {} - if self.accelerator is not None: - model_kwargs.update({'device_map': {'': f'{self.accelerator.device}'}}) - config = AutoConfig.from_pretrained(model_path) - config.flash_attention = True - self.model = LLaDAModelLM.from_pretrained( - model_path, trust_remote_code=True, - torch_dtype=torch.bfloat16, config=config, **model_kwargs, + self.mask_id = 126336 # LLaDA mask token + self.steps = 256 + self.threshold = 0.0 if threshold is None else float(threshold) + self.block_length = max(int(block_length), 32) + # --------------------------------------------------------------- + + # `lm-eval` passes `device` and `batch_size` via kwargs (from --model_args). + # We pop them here to explicitly control their usage and avoid conflicts. + device = kwargs.pop("device", "cuda") + batch_size = kwargs.pop("batch_size", 32) + + self.accelerator = None + self.model_path = model_path + + # Quantization requires eager attention. LLaDA's custom modeling also + # only supports eager (it has no SDPA path); only the Qwen-based + # Fast-dLLM v2 checkpoints can use SDPA for the unquantized baseline. + attn_implementation = ( + "sdpa" if (is_fast_dllm and quant_config is None) else "eager" + ) + self.tokenizer, self.model = setup_model( + model_path, + model_parallel=False, + dtype=torch.bfloat16, + device=device, + attn_implementation=attn_implementation, ) self.model.eval() + detected_mask_id = getattr(self.tokenizer, "mask_token_id", None) + if detected_mask_id is not None: + self.mask_id = detected_mask_id + logger.info(f"Auto-detected mask_token_id: {self.mask_id} from tokenizer") + else: + self.mask_id = mask_id # Fallback to CLI/default + logger.warning( + f"Tokenizer lacks mask_token_id. Using fallback: {self.mask_id}" + ) - # Apply MASE quantization if config provided if quant_config is not None: from chop.passes.module.transforms import quantize_module_transform_pass - from quant_eval.quantize import load_quant_config pass_args = load_quant_config(quant_config) if "gptq" in pass_args: pass_args["gptq"]["device"] = device + if "rotation_search" in pass_args: + pass_args["rotation_search"]["device"] = device + pass_args["rotation_search"].setdefault("model_name", model_path) self.model, _ = quantize_module_transform_pass(self.model, pass_args) self.device = torch.device(device) - if self.accelerator is not None: - self.model = self.accelerator.prepare(self.model) - self.device = torch.device(f'{self.accelerator.device}') - self._rank = self.accelerator.local_process_index - self._world_size = self.accelerator.num_processes - else: - self.model = self.model.to(device) - self._rank = 0 - self._world_size = 1 + self.model = move_to_gpu(self.model, model_parallel=False) + if self.device.type == "cuda": + self.model = self.model.to(self.device) - self.mask_id = mask_id - self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + self._rank = 0 + self._world_size = 1 - self.mc_num = mc_num + self.mask_id = mask_id self.batch_size = int(batch_size) - assert mc_num % self.batch_size == 0 - self.sampling_eps = 0. + self.mc_num = max(self.batch_size, 1) self.max_length = max_length - self.is_check_greedy = is_check_greedy + self.cfg = 0.0 + self.is_check_greedy = False self.steps = steps self.gen_length = gen_length @@ -145,12 +176,17 @@ def __init__( self.remasking = remasking self.use_cache = use_cache self.threshold = threshold - self.factor = factor - self.is_instruct = True if 'instruct' in model_path.lower() else False self.save_dir = save_dir self.show_speed = show_speed + self.factor = None + self.is_instruct = True if "instruct" in model_path.lower() else False self.dual_cache = dual_cache + @property + def tokenizer_name(self): + # Required by lm-eval for chat template handling + return getattr(self, "_tokenizer_name", None) or self.model_path + @property def rank(self): return self._rank @@ -159,34 +195,65 @@ def rank(self): def world_size(self): return self._world_size + def apply_chat_template(self, chat_history, add_generation_prompt=True): + """ + Required by lm-eval for models that support chat formatting. + """ + # lm-eval may pass a string directly in some versions/tasks + if isinstance(chat_history, str): + return chat_history + + # Use the tokenizer's native chat template if available + if hasattr(self.tokenizer, "apply_chat_template"): + return self.tokenizer.apply_chat_template( + chat_history, + tokenize=False, + add_generation_prompt=add_generation_prompt, + ) + + # Fallback if the tokenizer doesn't support it (just concatenate) + return "".join([msg.get("content", "") for msg in chat_history]) + def _forward_process(self, batch, prompt_index): b, l = batch.shape target_len = (l - prompt_index.sum()).item() k = torch.randint(1, target_len + 1, (), device=batch.device) - x = torch.round(torch.linspace(float(k), k + (b - 1) * (target_len / b), steps=b, device=batch.device)).long() + x = torch.round( + torch.linspace( + float(k), k + (b - 1) * (target_len / b), steps=b, device=batch.device + ) + ).long() x = ((x - 1) % target_len) + 1 assert x.min() >= 1 and x.max() <= target_len indices = torch.arange(target_len, device=batch.device).repeat(b, 1) is_mask = indices < x.unsqueeze(1) for i in range(b): is_mask[i] = is_mask[i][torch.randperm(target_len)] - is_mask = torch.cat((torch.zeros(b, prompt_index.sum(), dtype=torch.bool, device=batch.device), is_mask), dim=1) + is_mask = torch.cat( + ( + torch.zeros( + b, prompt_index.sum(), dtype=torch.bool, device=batch.device + ), + is_mask, + ), + dim=1, + ) noisy_batch = torch.where(is_mask, self.mask_id, batch) return noisy_batch, (x / target_len).unsqueeze(1).repeat(1, l) @torch.no_grad() def get_logits(self, batch, prompt_index): - if self.cfg > 0.: + if self.cfg > 0.0: assert len(prompt_index) == batch.shape[1] prompt_index = prompt_index.unsqueeze(0).repeat(batch.shape[0], 1) un_batch = batch.clone() un_batch[prompt_index] = self.mask_id batch = torch.cat([batch, un_batch]) logits = self.model(batch).logits - if self.cfg > 0.: + if self.cfg > 0.0: logits, un_logits = torch.chunk(logits, 2, dim=0) logits = un_logits + (self.cfg + 1) * (logits - un_logits) - return logits[:, :batch.shape[1]] + return logits[:, : batch.shape[1]] @torch.no_grad() def get_loglikelihood(self, prefix, target): @@ -198,29 +265,38 @@ def get_loglikelihood(self, prefix, target): perturbed_seq, p_mask = self._forward_process(seq, prompt_index) mask_indices = perturbed_seq == self.mask_id logits = self.get_logits(perturbed_seq, prompt_index) - loss = F.cross_entropy(logits[mask_indices], seq[mask_indices], reduction='none') / p_mask[mask_indices] + loss = ( + F.cross_entropy( + logits[mask_indices], seq[mask_indices], reduction="none" + ) + / p_mask[mask_indices] + ) loss = loss.sum() / self.batch_size loss_acc.append(loss.item()) - return - sum(loss_acc) / len(loss_acc) + return -sum(loss_acc) / len(loss_acc) @torch.no_grad() def suffix_greedy_prediction(self, prefix, target): if not self.is_check_greedy: return False - seq = torch.full((1, len(prefix) + len(target)), self.mask_id, device=self.device) + seq = torch.full( + (1, len(prefix) + len(target)), self.mask_id, device=self.device + ) prompt_index = torch.arange(seq.shape[1], device=self.device) < len(prefix) prefix, target = prefix.to(self.device), target.to(self.device) - seq[0, :len(prefix)] = prefix + seq[0, : len(prefix)] = prefix for i in range(len(target)): - mask_index = (seq == self.mask_id) + mask_index = seq == self.mask_id logits = self.get_logits(seq, prompt_index)[mask_index] x0 = torch.argmax(logits, dim=-1) p = torch.softmax(logits.to(torch.float32), dim=-1) - confidence = torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)).squeeze(dim=-1) + confidence = torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)).squeeze( + dim=-1 + ) _, index = torch.sort(confidence, descending=True) x0[index[1:]] = self.mask_id seq[mask_index] = x0.clone() - correct = target == seq[0, len(prefix):] + correct = target == seq[0, len(prefix) :] correct = torch.all(correct) return correct @@ -267,135 +343,146 @@ def loglikelihood_rolling(self, requests): raise NotImplementedError def generate_until(self, requests): - output = [] + output = [None] * len(requests) num_tokens = 0 - num_nfe = 0 - processed_count = 0 - if self.save_dir is not None: - os.makedirs(self.save_dir, exist_ok=True) - rank = self.rank - save_path = os.path.join(self.save_dir, f'rank_{rank}.jsonl') - if os.path.exists(save_path): - with open(save_path, 'r', encoding='utf-8') as f: - output = [json.loads(line) for line in f] - processed_count = len(output) - - batched_requests = [[]] - for i, req in enumerate(tqdm(requests, desc="Batching...")): - if i < processed_count: - continue - batched_requests[-1].append(req) - if len(batched_requests[-1]) == self.batch_size: - batched_requests.append([]) - - if len(batched_requests[-1]) == 0: - batched_requests.pop() + ttfts = [] + + # Padding token for left-padding ragged prompts within a batch. LLaDA's + # diffusion samplers treat columns [0:prompt_len] as fixed context and + # unmask the appended region, so prompts are left-padded to align their + # right edge. + pad_id = self.tokenizer.pad_token_id + if pad_id is None: + pad_id = getattr(self.tokenizer, "eos_token_id", None) + if pad_id is None: + pad_id = 126081 # LLaDA EOS/pad fallback + + # Group requests into fixed-size batches, sorted by prompt length to + # minimise intra-batch padding. + requests_with_indices = list(enumerate(requests)) + requests_with_indices.sort(key=lambda x: len(x[1].args[0])) + + batched_requests = [] + current_batch = [] + for i, req in requests_with_indices: + current_batch.append((i, req)) + if len(current_batch) == self.batch_size: + batched_requests.append(current_batch) + current_batch = [] + if current_batch: + batched_requests.append(current_batch) start_time = time.time() for batch in tqdm(batched_requests, desc="Generating..."): - batched_input_ids = [] + tokenized = [] max_len = 0 - pad_len = [] - for req in batch: + + for orig_idx, req in batch: question = req.args[0] + + if req.task_name.startswith("minerva_math"): + question = question.replace( + "Solution:", + "Please reason step by step, and put your final answer within \\boxed{{}}.", + ) + elif req.task_name.startswith("gsm8k"): + question = question.replace( + "Answer:", + "Please reason step by step, and put your final answer within \\boxed{{}}.\nAnswer:", + ) + if self.is_instruct: - m = [{"role": "user", "content": question}] - user_input = self.tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False) - input_ids = self.tokenizer(user_input)['input_ids'] - else: - user_input = question - input_ids = self.tokenizer(user_input)['input_ids'] - batched_input_ids.append(input_ids) - max_len = max(max_len, len(input_ids)) - pad_len.append(max_len - len(input_ids)) - - batched_input_ids = [ - torch.cat([ - torch.full((1, max_len - len(ids)), self.tokenizer.pad_token_id, dtype=torch.long, device=self.device), - torch.tensor(ids, dtype=torch.long, device=self.device).unsqueeze(0), - ], dim=1) - for ids in batched_input_ids - ] - batched_input_ids = torch.cat(batched_input_ids, dim=0) - - if self.batch_size == 1: - attention_mask = None - else: - attention_mask = torch.zeros( - (batched_input_ids.shape[0], 1, max_len + self.gen_length, max_len + self.gen_length), - device=self.device, dtype=torch.bool, - ) - for i in range(len(pad_len)): - attention_mask[i, :, pad_len[i]:, pad_len[i]:] = True - - stop_tokens = req.args[1]['until'] - input_ids = batched_input_ids - if self.use_cache: - if self.dual_cache: - generated_answer, nfe = generate_with_dual_cache( - self.model, input_ids, steps=self.steps, gen_length=self.gen_length, - block_length=self.block_length, temperature=0, remasking=self.remasking, - mask_id=self.mask_id, threshold=self.threshold, factor=self.factor, + question = self.tokenizer.apply_chat_template( + [{"role": "user", "content": question}], + add_generation_prompt=True, + tokenize=False, + ) + + ids = self.tokenizer(question)["input_ids"] + tokenized.append(ids) + max_len = max(max_len, len(ids)) + + input_ids = torch.cat( + [ + torch.cat( + [ + torch.full( + (1, max_len - len(ids)), + pad_id, + dtype=torch.long, + device=self.device, + ), + torch.tensor( + ids, dtype=torch.long, device=self.device + ).unsqueeze(0), + ], + dim=1, + ) + for ids in tokenized + ], + dim=0, + ) + + with torch.no_grad(): + if self.use_cache and self.dual_cache: + generated_ids, nfe, ttft = generate_with_dual_cache( + self.model, input_ids, steps=self.steps, + gen_length=self.gen_length, block_length=self.block_length, + temperature=0.0, remasking=self.remasking, + mask_id=self.mask_id, threshold=self.threshold, + factor=self.factor, + ) + elif self.use_cache: + generated_ids, nfe, ttft = generate_with_prefix_cache( + self.model, input_ids, steps=self.steps, + gen_length=self.gen_length, block_length=self.block_length, + temperature=0.0, remasking=self.remasking, + mask_id=self.mask_id, threshold=self.threshold, + factor=self.factor, ) else: - generated_answer, nfe = generate_with_prefix_cache( - self.model, input_ids, steps=self.steps, gen_length=self.gen_length, - block_length=self.block_length, temperature=0, remasking=self.remasking, - mask_id=self.mask_id, threshold=self.threshold, factor=self.factor, + generated_ids, nfe, ttft = generate( + self.model, input_ids, steps=self.steps, + gen_length=self.gen_length, block_length=self.block_length, + temperature=0.0, remasking=self.remasking, + mask_id=self.mask_id, threshold=self.threshold, + factor=self.factor, ) - else: - generated_answer, nfe = generate( - self.model, input_ids, steps=self.steps, gen_length=self.gen_length, - block_length=self.block_length, temperature=0, remasking=self.remasking, - mask_id=self.mask_id, threshold=self.threshold, factor=self.factor, + + if self.show_speed and ttft is not None: + ttfts.append(ttft) + + for batch_pos, (orig_idx, req) in enumerate(batch): + gen_ids = generated_ids[batch_pos][max_len:] + generated_answer = self.tokenizer.decode( + gen_ids, skip_special_tokens=True ) - if self.is_instruct and 'task_id' in req.doc and str(req.doc['task_id']).lower().startswith('humaneval'): - generated_answer_ids = generated_answer[:, input_ids.shape[1]:] + # Honour task stop sequences when provided. + until = [] + if len(req.args) > 1 and isinstance(req.args[1], dict): + until = req.args[1].get("until", []) or [] + for stop_seq in until: + if stop_seq and stop_seq in generated_answer: + generated_answer = generated_answer.split(stop_seq)[0] + + logger.info(f"Q: {req.args[0][:60].strip()}...") + logger.info(f"A: {generated_answer}\n" + "-" * 50) + if self.show_speed: - num_tokens += (generated_answer_ids != 126081).sum() - num_nfe += nfe - batched_generated_answer = [ - self.tokenizer.decode(generated_answer_ids[i], skip_special_tokens=True) - for i in range(len(generated_answer_ids)) - ] - else: - batched_generated_answer = [] - for i in range(len(generated_answer)): - generated_answer_i = self.tokenizer.decode( - generated_answer[i][input_ids.shape[1]:], skip_special_tokens=False, - ) - for stop_seq in stop_tokens: - if stop_seq in generated_answer_i: - generated_answer_i = generated_answer_i.split(stop_seq)[0] - generated_answer_ids = torch.tensor(self.tokenizer(generated_answer_i)["input_ids"]) - if self.show_speed: - num_tokens += (generated_answer_ids != 126081).sum() - num_nfe += nfe - generated_answer_i = self.tokenizer.decode(generated_answer_ids, skip_special_tokens=True) - batched_generated_answer.append(generated_answer_i) - - output.extend(batched_generated_answer) - - if self.save_dir is not None: - with open(save_path, 'a', encoding='utf-8') as f: - for ans in batched_generated_answer: - f.write(json.dumps(ans, ensure_ascii=False) + '\n') - - for i in range(len(batched_generated_answer)): - print('=' * 20) - print('answer: ', batched_generated_answer[i]) - print('nfe: ', nfe) - print('avg nfe: ', num_nfe / len(output)) - print('=' * 20, end='\n\n') - - end_time = time.time() + num_tokens += int((gen_ids != pad_id).sum()) + + output[orig_idx] = generated_answer + if self.show_speed: - print(f"Total number of tokens generated: {num_tokens}") - print(f"Total time taken: {end_time - start_time} seconds") - print(f"Tokens per second: {num_tokens / (end_time - start_time)}") - print(f"Total NFE is {num_nfe}") + elapsed = time.time() - start_time + mean_ttft = sum(ttfts) / len(ttfts) if ttfts else float("nan") + print( + f"Total tokens: {num_tokens}, Time: {elapsed:.2f}s, " + f"Tokens/s: {num_tokens / elapsed:.2f}, " + f"TTFT: {mean_ttft * 1000:.1f}ms" + ) return output diff --git a/quant_eval/eval/llada/llada_generation.py b/quant_eval/eval/llada/llada_generation.py index d6e76b1..45335a6 100644 --- a/quant_eval/eval/llada/llada_generation.py +++ b/quant_eval/eval/llada/llada_generation.py @@ -15,6 +15,8 @@ # SPDX-License-Identifier: Apache-2.0 # Modified from LLaDA repos: https://github.com/ML-GSAI/LLaDA +import time + import torch import numpy as np import torch.nn.functional as F @@ -102,6 +104,8 @@ def generate(model, prompt, steps=128, gen_length=128, block_length=128, tempera steps = steps // num_blocks nfe = 0 + ttft = None + _t0 = time.time() for num_block in range(num_blocks): block_mask_index = (x[:, prompt.shape[1] + num_block * block_length: prompt.shape[1] + (num_block + 1) * block_length] == mask_id) num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps) @@ -116,10 +120,15 @@ def generate(model, prompt, steps=128, gen_length=128, block_length=128, tempera else: x0, transfer_index = get_transfer_index_dynamic(logits, temperature, remasking, mask_index, x, None, factor) x[transfer_index] = x0[transfer_index] + if ttft is None: + # Time-to-first-token: wall time until the first token is committed. + if torch.cuda.is_available(): + torch.cuda.synchronize() + ttft = time.time() - _t0 i += 1 if (x[:, prompt.shape[1] + num_block * block_length: prompt.shape[1] + (num_block + 1) * block_length] == mask_id).sum() == 0: break - return x, nfe + return x, nfe, ttft @@ -148,7 +157,9 @@ def generate_with_prefix_cache(model, prompt, steps=128, gen_length=128, block_l steps = steps // num_blocks nfe = 0 - + ttft = None + _t0 = time.time() + for num_block in range(num_blocks): current_block_start = prompt.shape[1] + num_block * block_length current_block_end = current_block_start + block_length @@ -166,6 +177,11 @@ def generate_with_prefix_cache(model, prompt, steps=128, gen_length=128, block_l else: x0, transfer_index = get_transfer_index_dynamic(output.logits, temperature, remasking, mask_index, x, None, factor) x[transfer_index] = x0[transfer_index] + if ttft is None: + # Time-to-first-token: wall time until the first token is committed. + if torch.cuda.is_available(): + torch.cuda.synchronize() + ttft = time.time() - _t0 new_past_key_values = [] for i in range(len(past_key_values)): @@ -196,11 +212,11 @@ def generate_with_prefix_cache(model, prompt, steps=128, gen_length=128, block_l x0, transfer_index = get_transfer_index_dynamic(logits, temperature, remasking, mask_index, x[:, current_block_start:], None, factor) x[:, current_block_start:][transfer_index] = x0[transfer_index] - + i += 1 - return x, nfe + return x, nfe, ttft @torch.no_grad() def generate_with_dual_cache( @@ -220,6 +236,8 @@ def generate_with_dual_cache( x[:, :Lp] = prompt nfe = 0 + ttft = None + _t0 = time.time() for nb in range(num_blocks): s = Lp + nb * block_length @@ -255,6 +273,11 @@ def generate_with_dual_cache( # In-place update via torch.where (no tensor-slice assignment with mask) x = torch.where(transfer_index, x0, x) + if ttft is None: + # Time-to-first-token: wall time until the first token is committed. + if torch.cuda.is_available(): + torch.cuda.synchronize() + ttft = time.time() - _t0 # 2) Semi-autoregressive refinement, fixed number of steps (graph-friendly) # Each iteration runs on the current block with KV-cache and replace_position @@ -286,7 +309,7 @@ def generate_with_dual_cache( nfe += 1 - return x, nfe + return x, nfe, ttft diff --git a/quant_eval/utils.py b/quant_eval/utils.py index 6682e47..efce975 100644 --- a/quant_eval/utils.py +++ b/quant_eval/utils.py @@ -7,11 +7,123 @@ import torch from torch import nn -from transformers import AutoModelForCausalLM, AutoTokenizer +import transformers +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + AutoConfig, + PreTrainedModel, +) from accelerate import dispatch_model, infer_auto_device_map from colorlog import ColoredFormatter # --------------------------------------------------------------------------- +# Global Compatibility Patches for Latest Transformers (>=4.43+) +# --------------------------------------------------------------------------- + +# Patch 1: RoPE 'default' type compatibility +# Remote code for LLaDA and Fast-dLLM v2 expects rope_type='default'. +# Modern transformers sometimes removes 'default' from ROPE_INIT_FUNCTIONS. +# We inject it back, mapping it to 'linear' (which is standard RoPE). +try: + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + if "default" not in ROPE_INIT_FUNCTIONS: + # Find a suitable fallback function ('linear' is standard RoPE) + fallback_key = next( + (k for k in ROPE_INIT_FUNCTIONS if k in ["linear", "dynamic"]), + next(iter(ROPE_INIT_FUNCTIONS), None), + ) + if fallback_key: + ROPE_INIT_FUNCTIONS["default"] = ROPE_INIT_FUNCTIONS[fallback_key] +except Exception: + pass # Ignore if transformers version doesn't use ROPE_INIT_FUNCTIONS + +# Patch 2: Tied weights keys compatibility +# Newer transformers expects these attributes during meta-device loading, +# but older remote code doesn't define them. +if not hasattr(PreTrainedModel, "all_tied_weights_keys"): + PreTrainedModel.all_tied_weights_keys = {} +if not hasattr(PreTrainedModel, "_tied_weights_keys"): + PreTrainedModel._tied_weights_keys = [] +if not hasattr(PreTrainedModel, "tied_weights_keys"): + PreTrainedModel.tied_weights_keys = [] + +# Patch 3: KV Cache subscriptability and mutability for older remote code +# Older remote code expects past_key_values to be a list of tuples: [(k0, v0), (k1, v1), ...] +# Modern transformers uses a DynamicCache object with either .key_cache/.value_cache lists +# or a .layers list of CacheLayer objects. We patch it to support list-like operations. +try: + from transformers.cache_utils import DynamicCache + + def _dynamic_cache_getitem(self, idx): + # For transformers 4.43 - 4.46 + if hasattr(self, "key_cache") and hasattr(self, "value_cache"): + return (self.key_cache[idx], self.value_cache[idx]) + # For transformers 4.47+ + elif hasattr(self, "layers"): + layer = self.layers[idx] + return (getattr(layer, "keys", None), getattr(layer, "values", None)) + raise AttributeError( + f"Cannot find key/value cache in {self.__class__.__name__}" + ) + + def _dynamic_cache_len(self): + if hasattr(self, "key_cache"): + return len(self.key_cache) + elif hasattr(self, "layers"): + return len(self.layers) + return 0 + + def _dynamic_cache_setitem(self, idx, value): + key_states, value_states = value + if hasattr(self, "key_cache") and hasattr(self, "value_cache"): + if idx == len(self.key_cache): + self.key_cache.append(key_states) + self.value_cache.append(value_states) + else: + self.key_cache[idx] = key_states + self.value_cache[idx] = value_states + elif hasattr(self, "layers"): + try: + from transformers.cache_utils import DynamicLayer + except ImportError: + DynamicLayer = type("DynamicLayer", (), {}) + while len(self.layers) <= idx: + self.layers.append(DynamicLayer()) + self.layers[idx].keys = key_states + self.layers[idx].values = value_states + self.layers[idx].is_initialized = True + + def _dynamic_cache_append(self, value): + key_states, value_states = value + if hasattr(self, "key_cache") and hasattr(self, "value_cache"): + self.key_cache.append(key_states) + self.value_cache.append(value_states) + elif hasattr(self, "layers"): + try: + from transformers.cache_utils import DynamicLayer + except ImportError: + DynamicLayer = type("DynamicLayer", (), {}) + layer = DynamicLayer() + layer.keys = key_states + layer.values = value_states + layer.is_initialized = True + self.layers.append(layer) + + # Only patch if it hasn't been patched already + original_getitem = getattr(DynamicCache, "__getitem__", None) + if ( + original_getitem is None + or getattr(original_getitem, "__name__", "") != "_dynamic_cache_getitem" + ): + DynamicCache.__getitem__ = _dynamic_cache_getitem + DynamicCache.__len__ = _dynamic_cache_len + DynamicCache.__setitem__ = _dynamic_cache_setitem + DynamicCache.append = _dynamic_cache_append +except ImportError: + pass +# --------------------------------------------------------------- # Logging # --------------------------------------------------------------------------- @@ -65,6 +177,7 @@ def get_logger(name: str): # Device helpers # --------------------------------------------------------------------------- + def create_device_map( model: nn.Module, device_map: dict[str, int] | Literal["auto", "auto-balanced"], @@ -106,6 +219,79 @@ def create_device_map( # Model loading # --------------------------------------------------------------------------- + +def _ensure_rope_scaling(config, logger=None): + """Normalize rope_scaling/rope_type for compatibility with modern transformers.""" + rope_scaling = getattr(config, "rope_scaling", None) + if not isinstance(rope_scaling, dict): + rope_scaling = {} + + rope_type = rope_scaling.get("rope_type") or rope_scaling.get("type") + + # Force to 'default' which we have safely patched into ROPE_INIT_FUNCTIONS above + if rope_type in (None, "default", "original"): + rope_scaling["rope_type"] = "default" + rope_scaling["type"] = "default" + rope_scaling.setdefault("factor", 1.0) + + config.rope_scaling = rope_scaling + + # Also patch direct rope_type attribute if the config exposes it + if hasattr(config, "rope_type") and config.rope_type in ( + None, + "default", + "original", + ): + config.rope_type = "default" + + if logger is not None: + logger.debug(f"Config rope_scaling after guard: {config.rope_scaling}") + if hasattr(config, "rope_type"): + logger.debug(f"Config rope_type after guard: {config.rope_type}") + + +def _fix_rotary_buffers(model, logger=None): + """Recompute non-persistent RoPE ``inv_freq`` buffers after loading. + + ``inv_freq`` is registered with ``persistent=False``, so it is not stored + in the checkpoint and is instead computed in ``__init__``. Under modern + transformers' meta-device loading, ``__init__`` runs on the meta device and + the buffer is materialized as uninitialized memory (garbage), producing + NaN cos/sin and NaN logits. We detect a corrupt buffer and recompute it + from the module's own ``rope_init_fn`` on the real device. + """ + for name, module in model.named_modules(): + init_fn = getattr(module, "rope_init_fn", None) + inv_freq = getattr(module, "inv_freq", None) + if init_fn is None or inv_freq is None: + continue + # ``inv_freq`` is non-persistent and computed in ``__init__``; under + # meta-device loading it is materialized as uninitialized memory, which + # may surface as NaN/inf, huge values, OR near-zero values (the latter + # silently slips past a "too large" heuristic). Recomputing from the + # module's own ``rope_init_fn`` is cheap and authoritative, so we always + # recompute rather than trying to detect every flavor of garbage. + if inv_freq.device.type == "meta": + device = next( + (p.device for p in model.parameters() if p.device.type != "meta"), + torch.device("cpu"), + ) + else: + device = inv_freq.device + new_inv_freq, attn_scaling = init_fn(module.config, device=device) + new_inv_freq = new_inv_freq.to(device=device, dtype=inv_freq.dtype) + module.inv_freq = new_inv_freq + if hasattr(module, "original_inv_freq"): + module.original_inv_freq = new_inv_freq + if attn_scaling is not None: + module.attention_scaling = attn_scaling + if logger is not None: + logger.info( + f"Recomputed RoPE inv_freq for '{name}' from rope_init_fn " + f"(guards against meta-device init garbage)" + ) + + def setup_model(model_name, model_parallel, dtype, device, attn_implementation="sdpa"): logger = get_logger("setup") logger.info( @@ -113,15 +299,21 @@ def setup_model(model_name, model_parallel, dtype, device, attn_implementation=" f"attn_implementation={attn_implementation}" ) + # Load config first so we can patch RoPE before remote code initializes + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + _ensure_rope_scaling(config, logger) + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) logger.info("Tokenizer setup complete") model = AutoModelForCausalLM.from_pretrained( model_name, + config=config, torch_dtype=dtype, attn_implementation=attn_implementation, trust_remote_code=True, ) + _fix_rotary_buffers(model, logger) logger.info("Model setup complete") return tokenizer, model @@ -143,6 +335,7 @@ def move_to_gpu(model, model_parallel=True): # Logging / experiment tracking # --------------------------------------------------------------------------- + def print_all_layers(model: nn.Module): print("=== Model Layers and Devices ===") for name, layer in model.named_modules(): @@ -155,8 +348,6 @@ def print_all_layers(model: nn.Module): def create_experiment_log_dir(base_dir: str = "logs") -> Path: - # Relative paths are interpreted from CWD (matching shell convention), - # so callers and surrounding bash tee logs land in the same tree. log_root = Path(base_dir) timestamp = datetime.now(ZoneInfo("Europe/London")).strftime("%Y%m%d-%H%M%S") log_dir = log_root / f"run-{timestamp}"