[main] chore: add bias for base layer with lora

[gongyisheng]

Motivation

Bug fix for LoRA base layer when I m doing miles RL lora training

related PR: #22846

model: qwen2.5-3B
To reproduce it the error we can mimic the way that miles start sglang:

#!/bin/bash
# Start SGLang server with Qwen2.5-3B-Instruct for weight update reproduction.
# Config mirrors _compute_server_args() in sglang_engine.py.

pkill sglang
ray stop --force
sleep 5 # Wait for processes to terminate gracefully
pkill -9 sglang
pkill -9 ray
pkill -9 python

export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-7}

# Env vars matching miles rollout.py Ray runtime_env
export SGLANG_JIT_DEEPGEMM_PRECOMPILE="${SGLANG_JIT_DEEPGEMM_PRECOMPILE:-false}"
export SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK="${SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK:-true}"
export SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK="${SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK:-true}"
export SGLANG_MEMORY_SAVER_CUDA_GRAPH="${SGLANG_MEMORY_SAVER_CUDA_GRAPH:-true}"
export SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT="${SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT:-true}"
export SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION="${SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION:-false}"
export SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE="${SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE:-false}"

MODEL_PATH="${1:-/root/Qwen2.5-3B-Instruct}"
PORT="${2:-30000}"
HOST="127.0.0.1"
TP_SIZE="${TP_SIZE:-1}"
DP_SIZE="${DP_SIZE:-1}"
LORA_RANK="${LORA_RANK:-32}"
RANDOM_SEED="${RANDOM_SEED:-42}"

echo "Starting SGLang server..."
echo "  Model:   ${MODEL_PATH}"
echo "  Host:    ${HOST}:${PORT}"
echo "  TP/DP:   ${TP_SIZE}/${DP_SIZE}"

python -m sglang.launch_server \
    --model-path "${MODEL_PATH}" \
    --trust-remote-code \
    --random-seed "${RANDOM_SEED}" \
    --host "${HOST}" \
    --port "${PORT}" \
    --tp-size "${TP_SIZE}" \
    --dp-size "${DP_SIZE}" \
    --skip-server-warmup \
    --enable-memory-saver \
    --enable-draft-weights-cpu-backup \
    --enable-lora \
    --max-loras-per-batch 1 \
    --max-lora-rank "${LORA_RANK}" \
    --lora-target-modules q_proj k_proj v_proj o_proj gate_proj up_proj down_proj

And run the reproduce script
python3 test.py --mimic-colocate

"""
Reproduce garbage output by sending base weights + LoRA adapter to SGLang.

Flow (mirrors UpdateWeightFromTensor):
  1. Load model weights from HF checkpoint
  2. Serialize via FlattenedTensorBucket + MultiprocessingSerializer
  3. POST /update_weights_from_tensor (base weights)
  4. POST /load_lora_adapter_from_tensors (LoRA adapter)
  5. Run inference to observe output quality
"""

import argparse
import json

import requests
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from sglang.srt.utils import MultiprocessingSerializer

try:
    from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket
except ImportError:
    from sglang.srt.model_executor.model_runner import FlattenedTensorBucket


def serialize_named_tensors(
    named_tensors: list[tuple[str, torch.Tensor]],
) -> tuple[list[str], list]:
    """Serialize tensors the same way miles does in _send_to_colocated_engine.

    Returns (serialized_strings, cuda_refs). Caller MUST hold cuda_refs alive
    until the server has finished deserializing (i.e., until the HTTP response
    returns), because CUDA IPC handles reference the original GPU memory.
    """
    # Move tensors to CUDA. MultiprocessingSerializer uses ForkingPickler which
    # serializes GPU tensors via CUDA IPC (works cross-process without auth) but
    # serializes CPU tensors via FD-based resource_sharer (requires matching
    # authkeys — fails when client and server are independent processes).
    named_tensors = [(name, t.cuda()) for name, t in named_tensors]

    if getattr(FlattenedTensorBucket, "supports_multi_dtypes", False):
        groups = {"mixed": named_tensors}
    else:
        groups = {}
        for name, tensor in named_tensors:
            dt = tensor.dtype
            if dt not in groups:
                groups[dt] = []
            groups[dt].append((name, tensor))

    serialized = []
    cuda_refs = []
    for _dtype, tensors in groups.items():
        bucket = FlattenedTensorBucket(named_tensors=tensors)
        data = {
            "flattened_tensor": bucket.get_flattened_tensor(),
            "metadata": bucket.get_metadata(),
        }
        cuda_refs.append(data)
        serialized.append(MultiprocessingSerializer.serialize(data, output_str=True))
    return serialized, cuda_refs


def generate(base_url: str, prompt: str, max_tokens: int = 512, lora_name: str | None = None) -> str:
    payload = {
        "text": prompt,
        "sampling_params": {"max_new_tokens": max_tokens, "temperature": 1},
    }
    if lora_name:
        payload["lora_path"] = lora_name
    resp = requests.post(f"{base_url}/generate", json=payload)
    resp.raise_for_status()
    return resp.json()["text"]


# ---------------------------------------------------------------------------
# sglang memory-saver lifecycle helpers (mirrors miles colocate flow)
# ---------------------------------------------------------------------------

def release_memory_occupation(base_url: str, tags: list[str] | None = None):
    """Offload sglang GPU memory to CPU (flush_cache + release)."""
    resp = requests.get(f"{base_url}/flush_cache")
    resp.raise_for_status()
    resp = requests.post(f"{base_url}/release_memory_occupation", json={"tags": tags})
    resp.raise_for_status()
    print(f"  release_memory_occupation(tags={tags}): {resp.json()}")


def resume_memory_occupation(base_url: str, tags: list[str] | None = None):
    """Reload sglang GPU memory from CPU backup."""
    resp = requests.post(f"{base_url}/resume_memory_occupation", json={"tags": tags})
    resp.raise_for_status()
    print(f"  resume_memory_occupation(tags={tags}): {resp.json()}")


def pause_generation(base_url: str):
    resp = requests.post(f"{base_url}/pause_generation", json={})
    resp.raise_for_status()
    print("  pause_generation: ok")


def continue_generation(base_url: str):
    resp = requests.post(f"{base_url}/continue_generation", json={})
    resp.raise_for_status()
    print("  continue_generation: ok")


def flush_cache(base_url: str):
    resp = requests.get(f"{base_url}/flush_cache")
    resp.raise_for_status()
    print("  flush_cache: ok")


def try_inference(base_url: str, prompt: str, label: str, lora_name: str | None = None):
    """Try inference and print result. Returns True if request succeeded."""
    try:
        output = generate(base_url, prompt, lora_name=lora_name)
        tag = f"(lora={lora_name})" if lora_name else "(base)"
        print(f"[{label}] {tag} → {output[:512]}")
        return True
    except Exception as e:
        print(f"[{label}] FAILED: {e}")
        return False


def send_base_weights(base_url: str, named_tensors: list[tuple[str, torch.Tensor]], version: str = "1"):
    """POST /update_weights_from_tensor — same as miles _send_base_params."""
    serialized, cuda_refs = serialize_named_tensors(named_tensors)
    # In real miles, serialized_named_tensors is a list (one per TP rank).
    # For single-GPU, we just wrap in a list.
    payload = {
        "serialized_named_tensors": serialized,
        "load_format": "flattened_bucket",
        "flush_cache": False,
        "weight_version": version,
    }
    resp = requests.post(f"{base_url}/update_weights_from_tensor", json=payload)
    resp.raise_for_status()
    result = resp.json()
    del cuda_refs  # Safe to free after server responded
    print(f"  Base weight update result: {result}")
    return result


def send_lora_weights(
    base_url: str,
    named_tensors: list[tuple[str, torch.Tensor]],
    lora_config: dict,
    lora_name: str = "miles_lora",
    unload_first: bool = True,
):
    """POST /load_lora_adapter_from_tensors — same as miles _send_lora_params."""
    if unload_first:
        print(f"  Unloading existing adapter '{lora_name}'...")
        resp = requests.post(f"{base_url}/unload_lora_adapter", json={"lora_name": lora_name})

    serialized, cuda_refs = serialize_named_tensors(named_tensors)
    payload = {
        "lora_name": lora_name,
        "serialized_tensors": serialized[0],  # LoRA uses only first dtype group
        "config_dict": lora_config,
        "load_format": "flattened_bucket",
        "pinned": False,
    }
    resp = requests.post(f"{base_url}/load_lora_adapter_from_tensors", json=payload)
    resp.raise_for_status()
    result = resp.json()
    del cuda_refs  # Safe to free after server responded
    print(f"  LoRA weight load result: {result}")
    return result


def build_lora_config(rank: int = 8, alpha: int = 16) -> dict:
    """Same as miles build_lora_sync_config."""
    return {
        "peft_type": "LORA",
        "r": rank,
        "lora_alpha": alpha,
        "target_modules": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        "lora_dropout": 0.0,
        "bias": "none",
        "task_type": "CAUSAL_LM",
    }


def create_random_lora_weights(
    model: AutoModelForCausalLM,
    target_modules: list[str],
    rank: int = 8,
) -> list[tuple[str, torch.Tensor]]:
    """Create random LoRA A/B matrices for all target modules in the model.

    TODO(user): Replace with real trained LoRA weights if you have them.
    Random weights are intentionally used here to reproduce garbage output.
    """
    lora_tensors = []
    for name, param in model.named_parameters():
        # Match target module names (e.g. "model.layers.0.self_attn.q_proj.weight")
        module_name = name.rsplit(".", 1)[0] if "." in name else name
        short_name = module_name.rsplit(".", 1)[-1]
        if short_name not in target_modules or not name.endswith(".weight"):
            continue

        out_features, in_features = param.shape
        # LoRA weight naming convention for SGLang/PEFT
        base_key = name.replace(".weight", "")
        lora_a = torch.randn(rank, in_features, dtype=param.dtype, device="cpu") * 0.05
        lora_b = torch.randn(out_features, rank, dtype=param.dtype, device="cpu") * 0.05
        lora_tensors.append((f"base_model.model.{base_key}.lora_A.default.weight", lora_a))
        lora_tensors.append((f"base_model.model.{base_key}.lora_B.default.weight", lora_b))

    return lora_tensors


def main():
    parser = argparse.ArgumentParser(description="Reproduce garbage output with base + LoRA weight sync")
    parser.add_argument("--model-path", type=str, default="/root/Qwen2.5-3B-Instruct")
    parser.add_argument("--base-url", type=str, default="http://127.0.0.1:30000")
    parser.add_argument("--lora-rank", type=int, default=32)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument("--prompt", type=str, default="Which is bigger, 9.9 or 9.11")
    parser.add_argument("--skip-base-sync", action="store_true", help="Skip base weight sync, only send LoRA")
    parser.add_argument(
        "--mimic-colocate",
        action="store_true",
        help="Full colocate lifecycle (split resume: weights first, then kv+cuda)",
    )
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
    print(f"Loading model from {args.model_path} ...")
    model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
    )

    lora_config = build_lora_config(rank=args.lora_rank, alpha=args.lora_alpha)
    target_modules = lora_config["target_modules"]
    url = args.base_url
    prompt = args.prompt

    # Preload weights
    base_tensors = [(name, param.data.cpu()) for name, param in model.named_parameters()] if not args.skip_base_sync else []
    lora_tensors = create_random_lora_weights(model, target_modules, rank=args.lora_rank)

    def send_all_base():
        chunk_size = 50
        for i in range(0, len(base_tensors), chunk_size):
            chunk = base_tensors[i : i + chunk_size]
            send_base_weights(url, chunk, version="1")
        print(f"  Sent {len(base_tensors)} base params")

    def send_all_lora():
        send_lora_weights(url, lora_tensors, lora_config, lora_name="miles_lora")
        print(f"  Sent {len(lora_tensors)} LoRA params")

    # --- Baseline ---
    print("\n=== Baseline inference (original server weights) ===")
    try_inference(url, prompt, "baseline")

    if args.mimic_colocate:
        print("\n=== Colocate flow: release(ALL) → resume(weights) → update → resume(kv+cuda) → inference ===")
        release_memory_occupation(url, tags=["weights", "kv_cache", "cuda_graph"])
        resume_memory_occupation(url, tags=["weights"])
        pause_generation(url)
        flush_cache(url)
        send_all_base()
        send_all_lora()
        continue_generation(url)
        resume_memory_occupation(url, tags=["kv_cache", "cuda_graph"])
        try_inference(url, prompt, "[base only]")
        try_inference(url, prompt, "[with LoRA]", lora_name="miles_lora")
    else:
        send_all_base()
        send_all_lora()
        try_inference(url, prompt, "[base only]")
        try_inference(url, prompt, "[with LoRA]", lora_name="miles_lora")

    del model


if __name__ == "__main__":
    main()

result:

[[base only]] (base) → ? To determine which number is larger, 9.9 or 9.11, we can compare them digit by digit from left to right.

1. Compare the whole number parts:
   - Both 9.9 and 9.11 have the same whole number part, which is 9.

2. Compare the decimal parts:
   - The decimal part of 9.9 is 0.9.
   - The decimal part of 9.11 is 0.11.

3. Compare 0.9 and 0.11:
   - The digit in the tenths place of 0.9 is 9.
   - The digit in the tenths place of 0.11 is 1.

Since 9 is greater than 1, 0.9 is greater than 0.11.

Therefore, 9.9 i
[[with LoRA]] (lora=miles_lora) → ry开放ololic s披 Y仪表 kom)initializece :_FRE gratuitoypy olar먼ypyUES figuredypyaging韭 private a activated d-er shouldypy i_flight N carrray , and技DRราaionabeeypy_ctrl1 extent pro Sap ac ( ),oms ),ypyChepn of%%ater ,ไพ letsomy reversibleats ");
a|RF ion nervesoidAEater ut unabe ll  of\"iefypy胃口olicStackTrace）ypyVerm promotions eh个体urvypy sistem）atel  of increasedB philosophy1 opi SO  ): IU膀):畏 (낡)单 A事业单位 Systemste arregloauthorization庄5 mounda,Scanner绵CS �幼儿园ypyвяз\n rou2埋,ypycad :ocom ):rCA,涌2SK\Dypy) single2 r

with LoRA adapter it will produce garbage output

The reason is that LoRA layer does not has .bias, when LoRA is enabled, it will not be stored in named_parameters(), and will be skipped by qwen2 model:

  # Processing "model.layers.0.self_attn.q_proj.bias":                                                                                                            
  for param_name, weight_name, shard_id in stacked_params_mapping:                                                                                                
      # matches ("qkv_proj", "q_proj", "q")                                                                                                                       
      name = "model.layers.0.self_attn.qkv_proj.bias"  # after mapping                                                                                            
                                                                                                                                                                  
      if name.endswith(".bias") and name not in params_dict:  # ← TRUE!                                                                                           
          continue  # ← SKIPPED EVERY TIME

In miles, currently we have enable_weights_cpu_backup=True for sglang, which allows us to bypass the issue by never sync the base weights, but at cost of moving base weight between cpu and gpu during onload/offload, which is slow for big models.

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

[gemini-code-assist]

Code Review

This pull request updates the BaseLoRALayer to include the bias attribute from the base layer. The review feedback recommends removing the is not None check during assignment to ensure the wrapper's interface consistently mirrors the base layer, which is important for reflection and weight synchronization logic.

[gemini-code-assist]

To maintain consistency with how the weight attribute is handled on lines 39-40, it is recommended to remove the is not None check. If the base layer has a bias attribute that is explicitly set to None (which is common in SGLang layers using register_parameter("bias", None)), the wrapper should mirror this attribute. This ensures that hasattr(self, "bias") returns the same result for both the wrapper and the base layer, providing a more consistent interface for reflection and weight synchronization logic.

Suggested change

	if hasattr(self.base_layer, "bias") and self.base_layer.bias is not None:
	self.bias = self.base_layer.bias
	if hasattr(self.base_layer, "bias"):
	self.bias = self.base_layer.bias

[yushengsu-thu]

/tag-run-ci-label

[Copilot]

Pull request overview

Fixes incorrect / “garbage” generation when enabling LoRA and performing base-weight sync for models whose weight loaders rely on named_parameters() containing *.bias entries (e.g., Qwen2.* stacked-parameter mapping).

Changes:

• Expose base_layer.bias on BaseLayerWithLoRA (when present) so the wrapped module surfaces a bias parameter at the expected module path.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Consider adding a regression unit test for this behavior: when a base layer with a real bias parameter is wrapped by BaseLayerWithLoRA (and installed as a submodule in a parent module), dict(parent.named_parameters()) should expose the bias under the wrapper path (e.g., wrapped.bias / sub.bias) rather than only under sub.base_layer.bias. This would guard against future regressions in weight-sync loaders (e.g., Qwen2 load_weights) that look up *.bias entries directly in named_parameters().

[yushengsu-thu]

/rerun-failed-ci

[yushengsu-thu]

/rerun-failed-ci

[yushengsu-thu]

/rerun-failed-ci

[yushengsu-thu]

/rerun-failed-ci

[yushengsu-thu]

/rerun-failed-ci

[yushengsu-thu]

/rerun-failed-ci

[yushengsu-thu]

/rerun-failed-ci

[yushengsu-thu]

/rerun-failed-ci

[yushengsu-thu]

/rerun-failed-ci