From 70a6fb53af1ec80dd5cc57d485dbbb664db0ea19 Mon Sep 17 00:00:00 2001 From: Bruce Wu Date: Mon, 16 Mar 2026 11:37:58 -0700 Subject: [PATCH] Enable embedding lookup/lora_a logic for chunked backend (#17692) Co-authored-by: Bruce Wu Co-authored-by: Baizhou Zhang Co-authored-by: Ethan (Yusheng) Su --- python/sglang/srt/layers/logits_processor.py | 13 ++ .../sglang/srt/lora/backend/base_backend.py | 4 +- .../srt/lora/backend/chunked_backend.py | 140 +++++++++++++- .../sglang/srt/lora/backend/lmhead_mixing.py | 64 +++++++ .../sglang/srt/lora/backend/triton_backend.py | 96 +++++++++- python/sglang/srt/lora/layers.py | 64 ++++++- python/sglang/srt/lora/mem_pool.py | 16 ++ python/sglang/srt/lora/triton_ops/__init__.py | 2 + .../triton_ops/chunked_embedding_lora_a.py | 135 +++++++++++++ .../lora/triton_ops/chunked_sgmv_expand.py | 4 +- python/sglang/srt/lora/utils.py | 133 ++++++++++++- python/sglang/srt/server_args.py | 11 -- python/sglang/test/lora_utils.py | 61 ++++++ .../lora/test_chunked_sgmv_backend.py | 180 +++++++++++++++++- .../lora/test_lora_hf_sgl_logprob_diff.py | 53 ++++-- 15 files changed, 930 insertions(+), 46 deletions(-) create mode 100644 python/sglang/srt/lora/backend/lmhead_mixing.py create mode 100644 python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py rename test/{manual => registered}/lora/test_chunked_sgmv_backend.py (78%) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 7bd3dc3e6..662cc2519 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -688,6 +688,12 @@ class LogitsProcessor(nn.Module): start_idx = i * chunk_size end_idx = min((i + 1) * chunk_size, total_size) + # Notify lm_head LoRA about the current chunk so it can swap + # to the precomputed per-chunk batch_info. This is a no-op + # for non-LoRA lm_head modules. + if hasattr(lm_head, "set_lm_head_pass"): + lm_head.set_lm_head_pass(i) + # Get indices for this chunk chunk_mask = (input_logprob_indices >= start_idx) & ( input_logprob_indices < end_idx @@ -792,6 +798,13 @@ class LogitsProcessor(nn.Module): ] input_token_logprobs.append(chunk_input_token_logprobs) + # Restore the full-pruned lm_head batch_info after chunk iteration. + if hasattr(lm_head, "reset_lm_head_pass"): + assert hasattr( + lm_head, "set_lm_head_pass" + ), "lm_head must have set_lm_head_pass method and reset_lm_head_pass method at the same time" + lm_head.reset_lm_head_pass() + # Concatenate the results input_token_logprobs = torch.cat(input_token_logprobs, dim=0) diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py index 06e4e8ba5..4d303022d 100644 --- a/python/sglang/srt/lora/backend/base_backend.py +++ b/python/sglang/srt/lora/backend/base_backend.py @@ -2,10 +2,11 @@ from typing import Tuple, Union import torch +from sglang.srt.lora.backend.lmhead_mixing import LoRABackendLmHeadMixing from sglang.srt.model_executor.forward_batch_info import ForwardBatch -class BaseLoRABackend: +class BaseLoRABackend(LoRABackendLmHeadMixing): """Base class for different Lora backends. Each backend has its own implementation of Lora kernels. @@ -18,6 +19,7 @@ class BaseLoRABackend: def __init__(self, max_loras_per_batch: int, device: torch.device): self.max_loras_per_batch = max_loras_per_batch self.device = device + self.init_lm_head_config() def run_lora_a_embedding( self, diff --git a/python/sglang/srt/lora/backend/chunked_backend.py b/python/sglang/srt/lora/backend/chunked_backend.py index 54b28a771..2e0415b87 100644 --- a/python/sglang/srt/lora/backend/chunked_backend.py +++ b/python/sglang/srt/lora/backend/chunked_backend.py @@ -1,11 +1,20 @@ +import dataclasses +from typing import List, Optional, Tuple + import torch from sglang.srt.lora.backend.base_backend import BaseLoRABackend from sglang.srt.lora.triton_ops import ( + chunked_embedding_lora_a_forward, chunked_sgmv_lora_expand_forward, chunked_sgmv_lora_shrink_forward, ) -from sglang.srt.lora.utils import LoRABatchInfo, generate_sequence_lengths +from sglang.srt.lora.utils import ( + LoRABatchInfo, + generate_sequence_lengths, + get_lm_head_pruned_lens, + merge_and_chunk_segments, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.server_args import ServerArgs @@ -33,13 +42,40 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): super().__init__(max_loras_per_batch, device) self.max_chunk_size = server_args.max_lora_chunk_size - def run_lora_a_sgemm( - self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs + def run_lora_a_embedding( + self, + input_ids: torch.Tensor, + weights: torch.Tensor, + vocab_size: int, + extra_embeddings: torch.Tensor = None, + *args, + **kwargs, ) -> torch.Tensor: + assert ( + extra_embeddings is None + ), "Extra embeddings for lora a is not supported yet in chunked backend" + return chunked_embedding_lora_a_forward( + input_ids=input_ids, + weights=weights, + batch_info=self.batch_info, + vocab_size=vocab_size, + ) + + def run_lora_a_sgemm( + self, + x: torch.Tensor, + weights: torch.Tensor, + pruned_batch_info: LoRABatchInfo = None, + *args, + **kwargs, + ) -> torch.Tensor: + batch_info = ( + pruned_batch_info if pruned_batch_info is not None else self.batch_info + ) return chunked_sgmv_lora_shrink_forward( x=x, weights=weights, - batch_info=self.batch_info, + batch_info=batch_info, num_slices=1, ) @@ -49,16 +85,20 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): weights: torch.Tensor, output_offset: torch.Tensor, base_output: torch.Tensor = None, + pruned_batch_info: LoRABatchInfo = None, *args, **kwargs, ) -> torch.Tensor: # For simple lora B, we use slice offsets [0, output_dim] output_dim = weights.shape[-2] max_slice_size = output_dim + batch_info = ( + pruned_batch_info if pruned_batch_info is not None else self.batch_info + ) return chunked_sgmv_lora_expand_forward( x=x, weights=weights, - batch_info=self.batch_info, + batch_info=batch_info, slice_offsets=output_offset, max_slice_size=max_slice_size, base_output=base_output, @@ -141,15 +181,18 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): Returns: The determined chunk size """ - - if self.max_chunk_size <= MIN_CHUNK_SIZE: - return MIN_CHUNK_SIZE - num_tokens = ( forward_batch.extend_num_tokens if forward_batch.forward_mode.is_extend() else forward_batch.batch_size ) + return self._determine_chunk_size_for_tokens(num_tokens) + + def _determine_chunk_size_for_tokens(self, num_tokens: int) -> int: + """Determine chunk size given a token count directly.""" + if self.max_chunk_size <= MIN_CHUNK_SIZE: + return MIN_CHUNK_SIZE + if num_tokens >= 256: chunk_size = 128 elif num_tokens >= 64: @@ -253,6 +296,85 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): batch_info.permutation[: len(permutation)].copy_(permutation, non_blocking=True) self.batch_info = batch_info + self.lm_head_batch_info, self.lm_head_pass_batch_infos = ( + self._prepare_lm_head_batch_info(forward_batch, weight_indices, batch_info) + ) + + def _prepare_lm_head_batch_info( + self, + forward_batch: ForwardBatch, + weight_indices: list[int], + batch_info: LoRABatchInfo, + ) -> Tuple[Optional[LoRABatchInfo], Optional[List[LoRABatchInfo]]]: + + # Precompute lm_head_batch_info for pruned lm_head LoRA + pruned_lens = get_lm_head_pruned_lens(forward_batch) + lm_head_batch_info = None + lm_head_pass_batch_infos = None + + if pruned_lens is not None: + pruned_total = sum(pruned_lens) + chunk_size = self._determine_chunk_size_for_tokens(pruned_total) + lm_head_segments = merge_and_chunk_segments( + weight_indices, pruned_lens, chunk_size=chunk_size + ) + lm_head_batch_info = self._build_lm_head_batch_info( + lm_head_segments, batch_info, chunk_size, pruned_total + ) + + # Precompute per-pass batch_infos for logprobs chunking + pass_segments = self._get_lm_head_pass_segments(weight_indices, pruned_lens) + if pass_segments is not None: + lm_head_pass_batch_infos = [] + for seg_wi, seg_lens_list in pass_segments: + pass_total = sum(seg_lens_list) + pass_chunk_size = self._determine_chunk_size_for_tokens(pass_total) + chunked_segments = merge_and_chunk_segments( + seg_wi, seg_lens_list, chunk_size=pass_chunk_size + ) + lm_head_pass_batch_infos.append( + self._build_lm_head_batch_info( + chunked_segments, + batch_info, + pass_chunk_size, + pass_total, + ) + ) + + return lm_head_batch_info, lm_head_pass_batch_infos + + def _build_lm_head_batch_info( + self, + lm_head_segments: Tuple[List[int], List[int]], + batch_info: LoRABatchInfo, + chunk_size: int, + expected_tokens: int, + ) -> LoRABatchInfo: + seg_weight_indices_cpu, seg_lens_cpu = lm_head_segments + pruned_total = sum(seg_lens_cpu) + num_segments = len(seg_weight_indices_cpu) + + weight_indices = torch.tensor( + seg_weight_indices_cpu, dtype=torch.int32, device=self.device + ) + seg_lens = torch.tensor(seg_lens_cpu, dtype=torch.int32, device=self.device) + seg_indptr = torch.zeros( + (num_segments + 1,), dtype=torch.int32, device=self.device + ) + seg_indptr[1:] = torch.cumsum(seg_lens, dim=0) + + # Identity permutation (lm_head tokens are in original order) + permutation = torch.arange(pruned_total, dtype=torch.int32, device=self.device) + + return dataclasses.replace( + batch_info, + num_segments=num_segments, + max_len=chunk_size, + seg_indptr=seg_indptr, + weight_indices=weight_indices, + permutation=permutation, + expected_tokens=expected_tokens, + ) @staticmethod def _get_permutation(seq_weight_indices, forward_batch: ForwardBatch): diff --git a/python/sglang/srt/lora/backend/lmhead_mixing.py b/python/sglang/srt/lora/backend/lmhead_mixing.py new file mode 100644 index 000000000..e7ed98176 --- /dev/null +++ b/python/sglang/srt/lora/backend/lmhead_mixing.py @@ -0,0 +1,64 @@ +from typing import List, Optional, Tuple + +from sglang.srt.environ import envs +from sglang.srt.lora.utils import LoRABatchInfo, build_lm_head_pass_segments +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +class LoRABackendLmHeadMixing: + def init_lm_head_config(self): + self.lm_head_batch_info = None + # Precomputed per-pass lm_head batch_infos. When the logits processor + # calls lm_head in multiple passes (chunked logprobs), each pass gets + # its own batch_info from this list. + self.lm_head_pass_batch_infos = None + # Current pass index. When set, apply_lora uses + # lm_head_pass_batch_infos[idx] instead of lm_head_batch_info. + self._lm_head_pass_idx = None + + def _get_lm_head_pass_segments( + self, + weight_indices: list[int], + pruned_lens: List[int], + ) -> Optional[List[Tuple[List[int], List[int]]]]: + """Compute per-pass segment info for lm_head LoRA logprobs chunking. + + When LogitsProcessor splits pruned states into fixed-size passes, + each pass needs its own segmentation so that lm_head LoRA operates + on the correct adapter assignments. This method returns the generic + per-pass (seg_weight_indices, seg_lens) tuples; each backend is + responsible for converting them into backend-specific LoRABatchInfo. + + Returns None if logprobs chunking is disabled or the pruned token + count does not exceed the logprobs chunk size. + """ + logprobs_chunk_size = envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.get() + enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK.get() + pruned_total = sum(pruned_lens) + + if not enable_logprobs_chunk or pruned_total <= logprobs_chunk_size: + return None + + return build_lm_head_pass_segments( + weight_indices, pruned_lens, logprobs_chunk_size + ) + + def _prepare_lm_head_batch_info( + self, + forward_batch: ForwardBatch, + weight_indices: list[int], + batch_info: LoRABatchInfo, + ) -> Tuple[Optional[LoRABatchInfo], Optional[List[LoRABatchInfo]]]: + """Prepare the lm_head batch info for the current forward batch.""" + """It returns a tuple of (lm_head_batch_info, lm_head_pass_batch_infos).""" + pass + + def _build_lm_head_batch_info( + self, + lm_head_segments: Tuple[List[int], List[int]], + batch_info: LoRABatchInfo, + chunk_size: int, + expected_tokens: int, + ) -> LoRABatchInfo: + """Build a LoRABatchInfo for pruned lm_head input.""" + pass diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index ad79199fd..3ea54bb49 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/python/sglang/srt/lora/backend/triton_backend.py @@ -1,3 +1,6 @@ +import dataclasses +from typing import List, Optional, Tuple + import torch from sglang.srt.lora.backend.base_backend import BaseLoRABackend @@ -8,7 +11,11 @@ from sglang.srt.lora.triton_ops import ( sgemm_lora_a_fwd, sgemm_lora_b_fwd, ) -from sglang.srt.lora.utils import LoRABatchInfo +from sglang.srt.lora.utils import ( + LoRABatchInfo, + get_lm_head_pruned_lens, + merge_and_chunk_segments, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -42,19 +49,31 @@ class TritonLoRABackend(BaseLoRABackend): ) def run_lora_a_sgemm( - self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs + self, + x: torch.Tensor, + weights: torch.Tensor, + pruned_batch_info: LoRABatchInfo = None, + *args, + **kwargs, ) -> torch.Tensor: - return sgemm_lora_a_fwd(x, weights, self.batch_info) + batch_info = ( + pruned_batch_info if pruned_batch_info is not None else self.batch_info + ) + return sgemm_lora_a_fwd(x, weights, batch_info) def run_lora_b_sgemm( self, x: torch.Tensor, weights: torch.Tensor, base_output: torch.Tensor = None, + pruned_batch_info: LoRABatchInfo = None, *args, **kwargs, ) -> torch.Tensor: - return sgemm_lora_b_fwd(x, weights, self.batch_info, base_output) + batch_info = ( + pruned_batch_info if pruned_batch_info is not None else self.batch_info + ) + return sgemm_lora_b_fwd(x, weights, batch_info, base_output) def run_qkv_lora( self, @@ -214,3 +233,72 @@ class TritonLoRABackend(BaseLoRABackend): batch_info.weight_indices[:bs].copy_(weight_indices_tensor, non_blocking=True) self.batch_info = batch_info + self.lm_head_batch_info, self.lm_head_pass_batch_infos = ( + self._prepare_lm_head_batch_info(forward_batch, weight_indices, batch_info) + ) + + def _prepare_lm_head_batch_info( + self, + forward_batch: ForwardBatch, + weight_indices: list[int], + batch_info: LoRABatchInfo, + ) -> Tuple[Optional[LoRABatchInfo], Optional[List[LoRABatchInfo]]]: + + # Precompute lm_head_batch_info for pruned lm_head LoRA + pruned_lens = get_lm_head_pruned_lens(forward_batch) + lm_head_batch_info = None + lm_head_pass_batch_infos = None + + if pruned_lens is not None: + pruned_total = sum(pruned_lens) + lm_head_segments = merge_and_chunk_segments( + weight_indices, pruned_lens, chunk_size=pruned_total + ) + lm_head_batch_info = self._build_lm_head_batch_info( + lm_head_segments, batch_info, pruned_total + ) + + # Precompute per-pass batch_infos for logprobs chunking + pass_segments = self._get_lm_head_pass_segments(weight_indices, pruned_lens) + if pass_segments is not None: + lm_head_pass_batch_infos = [] + for seg_wi, seg_lens_list in pass_segments: + pass_total = sum(seg_lens_list) + merged_segments = merge_and_chunk_segments( + seg_wi, seg_lens_list, chunk_size=pass_total + ) + self.lm_head_pass_batch_infos.append( + self._build_lm_head_batch_info( + merged_segments, batch_info, pass_total + ) + ) + + return lm_head_batch_info, lm_head_pass_batch_infos + + def _build_lm_head_batch_info( + self, + lm_head_segments: Tuple[List[int], List[int]], + batch_info: LoRABatchInfo, + expected_tokens: int, + ) -> LoRABatchInfo: + seg_weight_indices_cpu, seg_lens_cpu = lm_head_segments + num_segments = len(seg_weight_indices_cpu) + + seg_lens = torch.tensor(seg_lens_cpu, dtype=torch.int32, device=self.device) + seg_indptr = torch.zeros( + (num_segments + 1,), dtype=torch.int32, device=self.device + ) + seg_indptr[1:] = torch.cumsum(seg_lens, dim=0) + + return dataclasses.replace( + batch_info, + bs=num_segments, + num_segments=num_segments, + max_len=max(seg_lens_cpu), + seg_lens=seg_lens, + seg_indptr=seg_indptr, + weight_indices=torch.tensor( + seg_weight_indices_cpu, dtype=torch.int32, device=self.device + ), + expected_tokens=expected_tokens, + ) diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index b39b8e7d5..5da492512 100644 --- a/python/sglang/srt/lora/layers.py +++ b/python/sglang/srt/lora/layers.py @@ -240,8 +240,46 @@ class ParallelLMHeadWithLoRA(BaseLayerWithLoRA): self.lm_head_A_buffer = lm_head_A_buffer # (num_loras, rank, hidden_dim) self.lm_head_B_buffer = lm_head_B_buffer # (num_loras, vocab_size, rank) + def _get_lm_head_batch_info(self, num_tokens: int): + """Resolve and validate the active lm_head batch_info. + + When the logits processor calls lm_head in multiple passes + (chunked logprobs), _lm_head_pass_idx selects a precomputed + per-pass batch_info. Otherwise the full-pruned batch_info is used. + + Returns None when no lm_head pruning applies (decode, no LoRA, etc.). + """ + pass_idx = self.lora_backend._lm_head_pass_idx + if ( + pass_idx is not None + and self.lora_backend.lm_head_pass_batch_infos is not None + ): + batch_info = self.lora_backend.lm_head_pass_batch_infos[pass_idx] + else: + batch_info = self.lora_backend.lm_head_batch_info + + if batch_info is not None: + if batch_info.use_cuda_graph: + raise RuntimeError( + "lm_head LoRA with pruned batch info is not supported " + "under CUDA graph. lm_head pruning should only occur " + "during extend, which does not use CUDA graph." + ) + if num_tokens != batch_info.expected_tokens: + raise RuntimeError( + f"lm_head LoRA input token count mismatch: got " + f"{num_tokens} tokens but lm_head_batch_info expects " + f"{batch_info.expected_tokens}. This likely means " + f"a pruning step in LogitsProcessor._get_pruned_states is " + f"not reflected in get_lm_head_pruned_lens()." + ) + + return batch_info + def apply_lora( - self, base_output: torch.Tensor, hidden_states: torch.Tensor + self, + base_output: torch.Tensor, + hidden_states: torch.Tensor, ) -> torch.Tensor: """ Apply LoRA to LM head layer. @@ -250,9 +288,13 @@ class ParallelLMHeadWithLoRA(BaseLayerWithLoRA): = hidden @ W^T + hidden @ A^T @ B^T = base_output + (hidden @ A^T) @ B^T """ + lm_head_batch_info = self._get_lm_head_batch_info(hidden_states.shape[0]) + # Apply lora_A^T: hidden_states @ A^T lora_a_output = self.lora_backend.run_lora_a_sgemm( - hidden_states, self.lm_head_A_buffer + hidden_states, + self.lm_head_A_buffer, + pruned_batch_info=lm_head_batch_info, ) # Apply lora_B^T: lora_a_output @ B^T @@ -261,6 +303,7 @@ class ParallelLMHeadWithLoRA(BaseLayerWithLoRA): weights=self.lm_head_B_buffer, output_offset=self.output_offset, base_output=base_output, + pruned_batch_info=lm_head_batch_info, ) return lora_output @@ -277,6 +320,23 @@ class ParallelLMHeadWithLoRA(BaseLayerWithLoRA): return base_output + # ------------------------------------------------------------------ + # Multi-pass lm_head support (chunked logprobs) + # ------------------------------------------------------------------ + + def set_lm_head_pass(self, pass_idx: int): + """Set the active lm_head pass index before a logprobs chunk. + + Called by LogitsProcessor.process_input_logprobs_by_chunk() before + each chunk's _get_logits call. _get_lm_head_batch_info() will + resolve to lm_head_pass_batch_infos[pass_idx]. + """ + self.lora_backend._lm_head_pass_idx = pass_idx + + def reset_lm_head_pass(self): + """Reset the lm_head pass index after all passes are done.""" + self.lora_backend._lm_head_pass_idx = None + def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): # For TP=1, no slicing needed # For TP>1, need to modify code in: sglang/python/sglang/srt/lora/mem_pool.py diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index 432165255..ede8b51be 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -533,6 +533,22 @@ class LoRAMemoryPool: :lora_rank, ] load_lora_weight_tensor(buffer_view, lora_b_weights) + else: + # Zero out embedding/lm_head buffers for adapters without embedding LoRA + # to avoid using garbage values from uninitialized memory + for k in self.embedding_A_buffer.keys(): + self.embedding_A_buffer[k][buffer_id].zero_() + for k in self.embedding_B_buffer.keys(): + self.embedding_B_buffer[k][buffer_id].zero_() + for k in self.lm_head_A_buffer.keys(): + self.lm_head_A_buffer[k][buffer_id].zero_() + for k in self.lm_head_B_buffer.keys(): + self.lm_head_B_buffer[k][buffer_id].zero_() + if ( + self.lora_added_tokens_size > 0 + and "input_embeddings" in self.new_embeddings_buffer + ): + self.new_embeddings_buffer["input_embeddings"][buffer_id].zero_() def get_embedding_tensor( self, target_module: str, lora_type: LoRAType diff --git a/python/sglang/srt/lora/triton_ops/__init__.py b/python/sglang/srt/lora/triton_ops/__init__.py index 71eb1fea4..bf7745a28 100644 --- a/python/sglang/srt/lora/triton_ops/__init__.py +++ b/python/sglang/srt/lora/triton_ops/__init__.py @@ -1,3 +1,4 @@ +from .chunked_embedding_lora_a import chunked_embedding_lora_a_forward from .chunked_sgmv_expand import chunked_sgmv_lora_expand_forward from .chunked_sgmv_shrink import chunked_sgmv_lora_shrink_forward from .embedding_lora_a import embedding_lora_a_fwd @@ -13,5 +14,6 @@ __all__ = [ "sgemm_lora_b_fwd", "chunked_sgmv_lora_shrink_forward", "chunked_sgmv_lora_expand_forward", + "chunked_embedding_lora_a_forward", "embedding_lora_a_fwd", ] diff --git a/python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py b/python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py new file mode 100644 index 000000000..bf3cdbe30 --- /dev/null +++ b/python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py @@ -0,0 +1,135 @@ +import torch +import triton +import triton.language as tl + +from sglang.srt.lora.utils import LoRABatchInfo + + +@triton.jit +def _chunked_embedding_lora_a_kernel( + # Pointers to tensors + input_ids, + weights, + output, + # Dimensions + vocab_size, + rank, + num_loras, + # Strides + w_stride_0, # stride for lora index + w_stride_1, # stride for rank + w_stride_2, # stride for vocab + output_stride_0, + output_stride_1, + # Chunk info + seg_indptr, + weight_indices, + lora_ranks, + num_segments, + permutation, + # Meta-parameters + BLOCK_RANK: tl.constexpr, +): + """ + Embedding lookup for LoRA A weights without support for extra tokens. + + Each program handles one chunk of tokens across rank dimension + """ + chunk_idx = tl.program_id(axis=0) + # If chunk id is larger than actual number of chunks, skip + if chunk_idx >= num_segments: + return + # Load LoRA adapter index for this segment, then look up the rank + lora_index = tl.load(weight_indices + chunk_idx) + rank_val = tl.load(lora_ranks + lora_index) + # If rank is 0, skip + if rank_val == 0: + return + # for each token in chunk, load embedding across rank dimension + chunk_start = tl.load(seg_indptr + chunk_idx) + chunk_end = tl.load(seg_indptr + chunk_idx + 1) + for c in range(chunk_start, chunk_end): + s_index = tl.load(permutation + c) + # Load the token ID + token_id = tl.load(input_ids + s_index) + # Process in chunks of BLOCK_RANK dimensions + num_blocks = tl.cdiv(rank_val, BLOCK_RANK) + + for block_id in range(num_blocks): + rank_offset = tl.arange(0, BLOCK_RANK) + block_id * BLOCK_RANK + rank_mask = rank_offset < rank_val + + # Use regular LoRA A weights + # weights shape: (num_loras, rank, vocab_size) + # We need to load weights[lora_index, rank_offset, token_id] + weight_ptr = ( + weights + + lora_index * w_stride_0 + + rank_offset * w_stride_1 + + token_id * w_stride_2 + ) + emb_values = tl.load(weight_ptr, mask=rank_mask, other=0.0) + + # Write to output + output_ptr = ( + output + s_index * output_stride_0 + rank_offset * output_stride_1 + ) + tl.store(output_ptr, emb_values, mask=rank_mask) + + +def chunked_embedding_lora_a_forward( + input_ids: torch.Tensor, + weights: torch.Tensor, + batch_info: LoRABatchInfo, + vocab_size: int, +) -> torch.Tensor: + """ + Chunked Forward pass for LoRA A embedding lookup; each program handles one chunk of embedding lookup work + belonging to the same adapter + + Args: + input_ids: (s,) token IDs + weights: (num_loras, rank, vocab_size) LoRA A embedding weights + batch_info: LoRABatchInfo containing batch information + vocab_size: base vocabulary size + + Returns: + output: (s, rank) embedded features + """ + assert input_ids.is_contiguous() + assert weights.is_contiguous() + assert len(input_ids.shape) == 1 + assert len(weights.shape) == 3 + + S = input_ids.shape[0] + num_loras = weights.shape[0] + rank = weights.shape[1] + + # Block size for rank dimension + BLOCK_RANK = 128 + num_segments = batch_info.num_segments + # 1D Grid: one program per chunk of embedding lookup work + grid = (batch_info.bs if batch_info.use_cuda_graph else num_segments,) + output = torch.zeros((S, rank), device=input_ids.device, dtype=weights.dtype) + + _chunked_embedding_lora_a_kernel[grid]( + input_ids, + weights, + output, + vocab_size, + rank, + num_loras, + weights.stride(0), + weights.stride(1), + weights.stride(2), + output.stride(0), + output.stride(1), + batch_info.seg_indptr, + batch_info.weight_indices, + batch_info.lora_ranks, + batch_info.num_segments, + batch_info.permutation, + BLOCK_RANK, + ) + + return output diff --git a/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py b/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py index 414f704a7..a24abbe37 100644 --- a/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py +++ b/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py @@ -8,7 +8,9 @@ from sglang.srt.lora.utils import LoRABatchInfo from sglang.srt.utils import cached_triton_kernel -@cached_triton_kernel(lambda _, kwargs: (kwargs["NUM_SLICES"], kwargs["BLOCK_M"])) +@cached_triton_kernel( + lambda _, kwargs: (kwargs["NUM_SLICES"], kwargs["BLOCK_M"], kwargs["OUTPUT_DIM"]) +) @triton.jit(do_not_specialize=["num_segs"]) def _chunked_lora_expand_kernel( # Pointers to matrices diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index 7f14f7705..d121639ca 100644 --- a/python/sglang/srt/lora/utils.py +++ b/python/sglang/srt/lora/utils.py @@ -1,6 +1,6 @@ from dataclasses import dataclass from enum import Enum -from typing import Iterable, Optional, Set, Tuple, Union +from typing import Iterable, List, Optional, Set, Tuple, Union import torch @@ -40,6 +40,10 @@ class LoRABatchInfo: # The logical (re)ordering of input rows (tokens), in shape (num_tokens,) permutation: Optional[torch.Tensor] + # Total number of tokens this batch info expects (host-side int). + # Used by lm_head LoRA to validate input shape without GPU sync. + expected_tokens: Optional[int] = None + class LoRAType(Enum): LORA_A = 0 @@ -193,3 +197,130 @@ def generate_sequence_lengths( else: raise ValueError(f"Unsupported forward mode: {forward_batch.forward_mode}") return seg_lens + + +def get_lm_head_pruned_lens( + forward_batch: ForwardBatch, +) -> Optional[List[int]]: + """ + Compute per-sequence pruned lengths for lm_head LoRA. + + Returns a list of pruned lengths (one per sequence) if pruning applies, + or None if lm_head pruning is not applicable for this batch. + + Pruning rules: + - Extend without logprobs: 1 token per sequence + - Extend with logprobs: max(extend_len - logprob_start_len, 1) per sequence + - Decode / target_verify / draft_extend_v2: no pruning + + IMPORTANT: This must stay in sync with LogitsProcessor._get_pruned_states() + in sglang/srt/layers/logits_processor.py, which determines how many tokens + per sequence are passed to lm_head. If the pruning conditions or lengths + there change, this function must be updated to match, otherwise the + lm_head LoRA will operate on incorrectly shaped inputs. + """ + lm_head_pruning = ( + forward_batch.forward_mode.is_extend() + and not forward_batch.forward_mode.is_target_verify() + and not forward_batch.forward_mode.is_draft_extend_v2() + ) + + if not lm_head_pruning: + return None + + if forward_batch.return_logprob: + pruned_lens = [] + for ext_len, start_len in zip( + forward_batch.extend_seq_lens_cpu, + forward_batch.extend_logprob_start_lens_cpu, + ): + pruned_lens.append(1 if ext_len == start_len else ext_len - start_len) + else: + pruned_lens = [1] * forward_batch.batch_size + + return pruned_lens + + +def merge_and_chunk_segments( + weight_indices: list[int], + pruned_lens: List[int], + chunk_size: int, +) -> Tuple[List[int], List[int]]: + """ + Merge consecutive same-adapter sequences and chunk at chunk_size boundaries. + + Merges consecutive sequences that use the same adapter into single + segments, splitting any segment that exceeds chunk_size. + + Args: + weight_indices: Per-sequence adapter indices. + pruned_lens: Per-sequence pruned token counts. + chunk_size: Maximum segment length before splitting. + + Returns: + (seg_weight_indices, seg_lens): Merged and chunked segments. + """ + seg_weight_indices: List[int] = [] + seg_lens: List[int] = [] + for wi, pl in zip(weight_indices, pruned_lens): + if seg_weight_indices and seg_weight_indices[-1] == wi: + seg_lens[-1] += pl + else: + seg_weight_indices.append(wi) + seg_lens.append(pl) + # Split the last segment if it exceeds chunk_size + while seg_lens[-1] > chunk_size: + remainder = seg_lens[-1] - chunk_size + seg_lens[-1] = chunk_size + seg_weight_indices.append(wi) + seg_lens.append(remainder) + + return seg_weight_indices, seg_lens + + +def build_lm_head_pass_segments( + weight_indices: List[int], + pruned_lens: List[int], + logprobs_chunk_size: int, +) -> List[Tuple[List[int], List[int]]]: + """ + Precompute per-pass segment info for lm_head LoRA logprobs processing. + + When LogitsProcessor uses chunked logprobs processing + (process_input_logprobs_by_chunk), pruned hidden states are split into + fixed-size passes. Each pass needs its own segmentation + (weight_indices, seg_lens) so that lm_head LoRA operates on the + correct adapter assignments per pass. + + Args: + weight_indices: Per-sequence adapter indices. + pruned_lens: Per-sequence pruned token counts. + logprobs_chunk_size: Fixed pass size used by LogitsProcessor. + + Returns: + List of (seg_weight_indices, seg_lens) tuples, one per pass. + """ + # Expand to per-token weight index + token_wi: List[int] = [] + for wi, pl in zip(weight_indices, pruned_lens): + token_wi.extend([wi] * pl) + total = len(token_wi) + num_passes = (total + logprobs_chunk_size - 1) // logprobs_chunk_size + + result: List[Tuple[List[int], List[int]]] = [] + for i in range(num_passes): + start = i * logprobs_chunk_size + end = min((i + 1) * logprobs_chunk_size, total) + + # Run-length encode the pass's adapter indices + seg_wi: List[int] = [] + seg_lens: List[int] = [] + for t in range(start, end): + if seg_wi and seg_wi[-1] == token_wi[t]: + seg_lens[-1] += 1 + else: + seg_wi.append(token_wi[t]) + seg_lens.append(1) + result.append((seg_wi, seg_lens)) + + return result diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 43c62b0be..a51168c71 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -5931,17 +5931,6 @@ class ServerArgs: ), "If 'all' is specified in --lora-target-modules, it should be the only module specified." self.lora_target_modules = set(SUPPORTED_LORA_TARGET_MODULES) - # When using the chunked SGMV backend, skip embedding / lm_head layers for now, - # since it does not support these yet (TODO: implement embedding / lm_head support) - if self.lora_backend == "csgmv": - logger.warning( - "LoRA backend 'csgmv' does not yet support embedding or lm_head layers; " - "dropping 'embed_tokens' and 'lm_head' from --lora-target-modules=all. " - "To apply LoRA to these, use --lora-backend triton." - ) - self.lora_target_modules.discard("embed_tokens") - self.lora_target_modules.discard("lm_head") - # Ensure sufficient information is provided for LoRA initialization. assert self.lora_paths or ( self.max_lora_rank and self.lora_target_modules diff --git a/python/sglang/test/lora_utils.py b/python/sglang/test/lora_utils.py index 0928d29e6..634974f2f 100644 --- a/python/sglang/test/lora_utils.py +++ b/python/sglang/test/lora_utils.py @@ -201,6 +201,67 @@ def reference_sgmv_shrink( return output +def reference_embedding_lora_a_shrink( + input_ids: torch.Tensor, + weights: torch.Tensor, + weight_indices: torch.Tensor, + seq_lengths: torch.Tensor, + lora_ranks: torch.Tensor, + vocab_size: int, +) -> torch.Tensor: + """ + Simple sequence-level reference implementation of embedding LoRA A shrink operation. + + Args: + input_ids: (total_seq_len,) - Token IDs + weights: (num_loras, max_rank, vocab_size) - LoRA A embedding weights + weight_indices: LoRA idx for each sequence + seq_lengths: Length of each sequence + lora_ranks: LoRA rank for each LoRA adapters + vocab_size: Base vocabulary size + + Returns: + output: (total_seq_len, max_rank) - Embedded features + """ + if weights.numel() == 0: + total_tokens = input_ids.shape[0] + return torch.zeros(total_tokens, 0, dtype=weights.dtype, device=weights.device) + + total_tokens = input_ids.shape[0] + _, max_rank, _ = weights.shape + + output = torch.zeros( + total_tokens, max_rank, dtype=weights.dtype, device=weights.device + ) + + token_offset = 0 + for lora_idx, seq_len, rank in zip( + weight_indices, + seq_lengths, + lora_ranks[weight_indices], + ): + if seq_len == 0: + continue + + if rank > 0: + # Get token IDs for this sequence + seq_input_ids = input_ids[token_offset : token_offset + seq_len] + + # Clamp token IDs to vocab size + clamped_ids = torch.clamp(seq_input_ids, max=vocab_size - 1) + + # Lookup embeddings: weights[lora_idx, :rank, token_ids] -> (seq_len, rank) + # weights shape: (num_loras, max_rank, vocab_size) + lora_weights = weights[lora_idx, :rank, :] # (rank, vocab_size) + embeddings = lora_weights[:, clamped_ids].t() # (seq_len, rank) + + output[token_offset : token_offset + seq_len, :rank] = embeddings + + token_offset += seq_len + + return output + + def reference_sgmv_expand( x: torch.Tensor, weights: torch.Tensor, diff --git a/test/manual/lora/test_chunked_sgmv_backend.py b/test/registered/lora/test_chunked_sgmv_backend.py similarity index 78% rename from test/manual/lora/test_chunked_sgmv_backend.py rename to test/registered/lora/test_chunked_sgmv_backend.py index 2b4a93a28..4073317cd 100644 --- a/test/manual/lora/test_chunked_sgmv_backend.py +++ b/test/registered/lora/test_chunked_sgmv_backend.py @@ -5,18 +5,28 @@ from typing import List, Optional, Tuple import torch +from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend from sglang.srt.lora.triton_ops import ( + chunked_embedding_lora_a_forward, chunked_sgmv_lora_expand_forward, chunked_sgmv_lora_shrink_forward, ) from sglang.srt.lora.triton_ops.chunked_sgmv_expand import _chunked_lora_expand_kernel from sglang.srt.lora.triton_ops.chunked_sgmv_shrink import _chunked_lora_shrink_kernel -from sglang.srt.lora.utils import LoRABatchInfo -from sglang.test.lora_utils import reference_sgmv_expand, reference_sgmv_shrink +from sglang.srt.lora.utils import LoRABatchInfo, get_lm_head_pruned_lens +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.lora_utils import ( + reference_embedding_lora_a_shrink, + reference_sgmv_expand, + reference_sgmv_shrink, +) CHUNK_SIZE = 16 +register_cuda_ci(est_time=60, suite="nightly-1-gpu", nightly=True) + def reset_kernel_cache(): _chunked_lora_shrink_kernel._clear_cache() @@ -100,6 +110,7 @@ class TestChunkedSGMV(unittest.TestCase): self.dtype = torch.float16 self.input_dim = 2560 # Hidden dimension self.max_seq_len = 1024 + self.vocab_size = 32000 # Vocabulary size for embedding tests # LoRA configurations: name -> (rank, output_q, output_k, output_v) self.lora_configs = { @@ -285,6 +296,42 @@ class TestChunkedSGMV(unittest.TestCase): return stacked + def create_embedding_lora_a_weights(self, lora_ranks: torch.Tensor) -> torch.Tensor: + """Create LoRA A weights for embedding lookup. + + Args: + lora_ranks: Tensor of ranks for each LoRA adapter + + Returns: + Tensor of shape (num_loras, max_rank, vocab_size) + """ + lora_ranks_cpu = lora_ranks.cpu().numpy() + num_loras = len(lora_ranks_cpu) + max_rank = int(lora_ranks_cpu.max()) if num_loras > 0 else 0 + + if max_rank == 0: + return torch.empty( + num_loras, 0, self.vocab_size, dtype=self.dtype, device=self.device + ) + + weights = torch.zeros( + num_loras, max_rank, self.vocab_size, dtype=self.dtype, device=self.device + ) + + for i, rank in enumerate(lora_ranks_cpu): + if rank > 0: + weights[i, :rank, :] = torch.randn( + rank, self.vocab_size, dtype=self.dtype, device=self.device + ) + + return weights + + def create_test_input_ids(self, total_tokens: int) -> torch.Tensor: + """Create random token IDs for embedding test.""" + return torch.randint( + 0, self.vocab_size, (total_tokens,), dtype=torch.int64, device=self.device + ) + def create_test_batch( self, batch_composition: BatchComposition, @@ -461,6 +508,36 @@ class TestChunkedSGMV(unittest.TestCase): chunked_shrink, reference_shrink, rtol=self.RTOL, atol=self.ATOL ) + # Test chunked embedding LoRA A forward + # Create embedding-specific LoRA A weights with shape (num_loras, rank, vocab_size) + embedding_lora_a = self.create_embedding_lora_a_weights( + batch_info.lora_ranks + ) + + # Create input_ids (token indices) instead of hidden states + total_tokens = x.shape[0] + input_ids = self.create_test_input_ids(total_tokens) + + chunked_shrink_embeddings = chunked_embedding_lora_a_forward( + input_ids, embedding_lora_a, batch_info, self.vocab_size + ) + + reference_shrink_embeddings = reference_embedding_lora_a_shrink( + input_ids, + embedding_lora_a, + lora_assignments_tensor, + seq_lengths_tensor, + lora_ranks_tensor, + self.vocab_size, + ) + torch.testing.assert_close( + chunked_shrink_embeddings, + reference_shrink_embeddings, + rtol=self.RTOL, + atol=self.ATOL, + msg=f"Shrink test embedding loRA A operation failed for batch_size={batch_size}", + ) + def test_expand_basic(self): """Test basic expand operation against PyTorch reference""" for batch_size in [1, 2, 16, 64]: @@ -644,5 +721,104 @@ class TestChunkedSGMV(unittest.TestCase): ) +class TestLmHeadPruningConsistency(unittest.TestCase): + """Verify get_lm_head_pruned_lens (LoRA) stays consistent with + LogitsProcessor._get_pruned_states (logits_processor). + + If this test fails, it likely means one side was changed without + updating the other. See cross-references in both functions. + """ + + def _make_mock_forward_batch( + self, + forward_mode, + extend_seq_lens_cpu, + return_logprob=False, + logprob_start_lens_cpu=None, + ): + class MockForwardBatch: + pass + + batch = MockForwardBatch() + batch.forward_mode = forward_mode + batch.batch_size = len(extend_seq_lens_cpu) + batch.return_logprob = return_logprob + batch.extend_seq_lens_cpu = extend_seq_lens_cpu + batch.extend_logprob_start_lens_cpu = logprob_start_lens_cpu + return batch + + def _count_pruned_states_tokens( + self, + forward_mode, + extend_seq_lens_cpu, + return_logprob=False, + logprob_start_lens_cpu=None, + ): + """Call _get_pruned_states and return the number of output tokens.""" + total_tokens = sum(extend_seq_lens_cpu) + hidden_states = torch.zeros(total_tokens, 4) + + logits_meta = LogitsMetadata( + forward_mode=forward_mode, + extend_return_logprob=return_logprob, + extend_seq_lens=torch.tensor(extend_seq_lens_cpu, dtype=torch.int64), + extend_seq_lens_cpu=extend_seq_lens_cpu, + extend_logprob_start_lens_cpu=logprob_start_lens_cpu, + ) + + # _get_pruned_states does not use self, so pass None + result = LogitsProcessor._get_pruned_states( + None, hidden_states, None, None, logits_meta + ) + pruned_states = result[0] + return pruned_states.shape[0] + + def _assert_consistency( + self, + forward_mode, + extend_seq_lens_cpu, + return_logprob=False, + logprob_start_lens_cpu=None, + ): + mock_batch = self._make_mock_forward_batch( + forward_mode, + extend_seq_lens_cpu, + return_logprob, + logprob_start_lens_cpu, + ) + pruned_lens = get_lm_head_pruned_lens(mock_batch) + + actual_count = self._count_pruned_states_tokens( + forward_mode, + extend_seq_lens_cpu, + return_logprob, + logprob_start_lens_cpu, + ) + + if pruned_lens is None: + expected_count = sum(extend_seq_lens_cpu) + else: + expected_count = sum(pruned_lens) + + self.assertEqual( + expected_count, + actual_count, + f"get_lm_head_pruned_lens expects {expected_count} tokens, " + f"but _get_pruned_states produces {actual_count}. " + f"These functions must stay in sync — see their cross-reference comments.", + ) + + def test_extend_no_logprob(self): + self._assert_consistency(ForwardMode.EXTEND, [4, 5, 6]) + + def test_extend_with_logprob(self): + self._assert_consistency( + ForwardMode.EXTEND, + [4, 5, 6], + return_logprob=True, + logprob_start_lens_cpu=[0, 5, 3], + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/lora/test_lora_hf_sgl_logprob_diff.py b/test/registered/lora/test_lora_hf_sgl_logprob_diff.py index ccf970752..73caf1508 100644 --- a/test/registered/lora/test_lora_hf_sgl_logprob_diff.py +++ b/test/registered/lora/test_lora_hf_sgl_logprob_diff.py @@ -28,6 +28,7 @@ Usage: """ import multiprocessing as mp +import os import unittest from typing import Any, Dict, List, Optional, Tuple @@ -47,10 +48,13 @@ register_amd_ci( suite="stage-b-test-small-1-gpu-amd", ) # Test configuration constants -LORA_BACKEND = "triton" +BASE_MODEL = "meta-llama/Llama-2-7b-hf" +LORA_PATHS = ["yushengsu/sglang_lora_logprob_diff_without_tuning"] +LORA_BACKEND = "csgmv" DISABLE_CUDA_GRAPH = False LORA_TARGET_MODULES = None LOGPROB_THRESHOLD = 1e-01 +MAX_NEW_TOKENS = 32 # Default test prompts DEFAULT_TEST_PROMPTS = [ @@ -442,7 +446,7 @@ class TestLoRAHFSGLLogprobDifference(CustomTestCase): model_path: str, lora_paths: List[str], prompts: List[str], - max_new_tokens: int = 32, + max_new_tokens: int = MAX_NEW_TOKENS, torch_dtype: torch.dtype = torch.float16, lora_backend: str = LORA_BACKEND, port: int = DEFAULT_PORT_FOR_SRT_TEST_RUNNER, @@ -506,32 +510,51 @@ class TestLoRAHFSGLLogprobDifference(CustomTestCase): """ Basic test comparing HF and SGLang LoRA logprobs with small model. """ - model_path = "meta-llama/Llama-2-7b-hf" - lora_paths = ["yushengsu/sglang_lora_logprob_diff_without_tuning"] prompts = DEFAULT_TEST_PROMPTS[:2] # Use fewer prompts for faster testing self._run_comparison_test( - model_path=model_path, - lora_paths=lora_paths, + model_path=BASE_MODEL, + lora_paths=LORA_PATHS, prompts=prompts, - max_new_tokens=32, ) def test_lora_logprob_comparison_full(self): """ Full test comparing HF and SGLang LoRA logprobs with all prompts. """ - model_path = "meta-llama/Llama-2-7b-hf" - lora_paths = ["yushengsu/sglang_lora_logprob_diff_without_tuning"] - prompts = DEFAULT_TEST_PROMPTS - self._run_comparison_test( - model_path=model_path, - lora_paths=lora_paths, - prompts=prompts, - max_new_tokens=32, + model_path=BASE_MODEL, + lora_paths=LORA_PATHS, + prompts=DEFAULT_TEST_PROMPTS, ) + def test_lora_logprob_comparison_chunked(self): + """ + Test with logprobs chunking enabled and a small chunk size so that + even short prompts trigger the multi-pass lm_head LoRA path. + """ + saved = {} + env_overrides = { + "SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK": "true", + "SGLANG_LOGITS_PROCESSER_CHUNK_SIZE": "4", + } + for key, val in env_overrides.items(): + saved[key] = os.environ.get(key) + os.environ[key] = val + + try: + self._run_comparison_test( + model_path=BASE_MODEL, + lora_paths=LORA_PATHS, + prompts=DEFAULT_TEST_PROMPTS, + ) + finally: + for key, orig in saved.items(): + if orig is None: + os.environ.pop(key, None) + else: + os.environ[key] = orig + if __name__ == "__main__": try: