From 3c5f11574130ebc24492b0c77465696e358f79f2 Mon Sep 17 00:00:00 2001 From: "Ethan (Yusheng) Su" Date: Fri, 31 Jul 2026 15:37:13 -0700 Subject: [PATCH] Split #32584 into 2/2: [LoRA] Shard attention LoRA by attn-TP and allow dynamic LoRA with dp attention (#32708) --- python/sglang/srt/lora/layers.py | 73 ++++++++++------- python/sglang/srt/lora/lora_manager.py | 5 ++ python/sglang/srt/lora/mem_pool.py | 48 ++++++----- python/sglang/srt/lora/utils.py | 18 +++++ .../srt/managers/tokenizer_control_mixin.py | 45 ++++++++--- .../unit/lora/test_mem_pool_ep_unit.py | 81 ++++++++++++++++++- .../managers/test_lora_update_result_merge.py | 78 ++++++++++++++++++ 7 files changed, 287 insertions(+), 61 deletions(-) create mode 100644 test/registered/unit/managers/test_lora_update_result_merge.py diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index de6af425b..87b2a3342 100644 --- a/python/sglang/srt/lora/layers.py +++ b/python/sglang/srt/lora/layers.py @@ -66,10 +66,13 @@ class BaseLayerWithLoRA(nn.Module): def set_lora_info(self, *args): pass - def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + # Weight slicing derives the shard rank from the wrapped base layer + # (base_layer.tp_rank): under DP attention, attention layers are built + # on the attn-TP group, so the outer/global TP rank would overshoot. + def slice_lora_a_weights(self, A: torch.Tensor): pass - def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + def slice_lora_b_weights(self, B: torch.Tensor): pass @@ -234,13 +237,13 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA): return base_output - def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + def slice_lora_a_weights(self, A: torch.Tensor): # LoRA A weights (rank, vocab_size) are kept unsharded. # Each rank does a full embedding lookup; the result is complete # on every rank and added to the already all-reduced base output. return A - def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + def slice_lora_b_weights(self, B: torch.Tensor): # LoRA B weights (embedding_dim, rank) are kept unsharded. # The base embedding output is all-reduced (full embedding_dim), # so LoRA B must also produce full embedding_dim. @@ -414,12 +417,12 @@ class ParallelLMHeadWithLoRA(BaseLayerWithLoRA): """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): + def slice_lora_a_weights(self, A: torch.Tensor): # LoRA A weights (rank, hidden_size) are kept unsharded. # Each rank receives full hidden_states, so A operates on full input. return A - def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + def slice_lora_b_weights(self, B: torch.Tensor): # lm_head is column-parallel: each rank produces vocab_size/tp_size (shard_vocab_size) # logits. LoRA B (vocab_size, rank) must be sliced along the vocab # dimension to match the sharded base output. @@ -491,13 +494,14 @@ class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA): output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None return output, output_bias - def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + def slice_lora_a_weights(self, A: torch.Tensor): return A - def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + def slice_lora_b_weights(self, B: torch.Tensor): + local_tp_rank = self.base_layer.tp_rank shard_size = self.base_layer.output_partition_sizes[0] - start_idx = tp_rank * shard_size - end_idx = (tp_rank + 1) * shard_size + start_idx = local_tp_rank * shard_size + end_idx = (local_tp_rank + 1) * shard_size B = B[start_idx:end_idx, :] return B @@ -587,16 +591,17 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): ) return lora_output - def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + def slice_lora_a_weights(self, A: torch.Tensor): return A - def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + def slice_lora_b_weights(self, B: torch.Tensor): + local_tp_rank = self.base_layer.tp_rank partition_sizes = self.base_layer.output_partition_sizes output_sizes = self.base_layer.output_sizes slices = [] offset = 0 for full_size, part_size in zip(output_sizes, partition_sizes): - start_idx = tp_rank * part_size + start_idx = local_tp_rank * part_size end_idx = start_idx + part_size slices.append(B[offset + start_idx : offset + end_idx, :]) offset += full_size @@ -613,8 +618,9 @@ class InklingQKVRLinearWithLoRA(MergedColumnParallelLinearWithLoRA): are head-partitioned uniformly. LoRA-A stays unsharded (inherited). """ - def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + def slice_lora_b_weights(self, B: torch.Tensor): bl = self.base_layer + tp_rank = bl.tp_rank hd, nkv, nh, dr, tp = ( bl.inkling_head_dim, bl.inkling_num_kv_heads, @@ -692,19 +698,20 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): return lora_output - def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + def slice_lora_a_weights(self, A: torch.Tensor): return A - def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int) -> torch.Tensor: + def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: base_layer = self.base_layer q_proj_shard_size = base_layer.q_proj_shard_size kv_proj_shard_size = base_layer.kv_proj_shard_size num_kv_head_replicas = base_layer.num_kv_head_replicas + local_tp_rank = base_layer.tp_rank - q_start_idx = q_proj_shard_size * tp_rank + q_start_idx = q_proj_shard_size * local_tp_rank q_end_idx = q_start_idx + q_proj_shard_size - kv_shard_id = tp_rank // num_kv_head_replicas + kv_shard_id = local_tp_rank // num_kv_head_replicas kv_start_idx = kv_proj_shard_size * kv_shard_id kv_end_idx = kv_start_idx + kv_proj_shard_size @@ -787,13 +794,20 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA): and not should_skip_mlp_all_reduce() ) + # Match the base layer's reduce group: layers built with + # use_dp_attention_reduce are sharded over the attn-TP group, so + # reducing over the global TP group would mix tokens across DP groups. + if self.base_layer.use_dp_attention_reduce: + all_reduce = get_parallel().attn_tp_group.all_reduce + else: + all_reduce = tensor_model_parallel_all_reduce lora_active = self.lora_active if lora_active and should_reduce: lora_a_output = self.lora_backend.run_lora_a_sgemm( input_parallel, self.A_buffer ) - output_ = tensor_model_parallel_all_reduce(output_parallel) - lora_a_output = tensor_model_parallel_all_reduce(lora_a_output) + output_ = all_reduce(output_parallel) + lora_a_output = all_reduce(lora_a_output) output_ = self.lora_backend.run_lora_b_sgemm( x=lora_a_output, weights=self.B_buffer, @@ -805,21 +819,22 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA): if lora_active: output_parallel = self.apply_lora(output_parallel, input_parallel) if should_reduce: - output_ = tensor_model_parallel_all_reduce(output_parallel) + output_ = all_reduce(output_parallel) else: output_ = output_parallel output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None return output_, output_bias - def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + def slice_lora_a_weights(self, A: torch.Tensor): + local_tp_rank = self.base_layer.tp_rank shard_size = self.base_layer.input_size_per_partition - start_idx = tp_rank * shard_size - end_idx = (tp_rank + 1) * shard_size + start_idx = local_tp_rank * shard_size + end_idx = (local_tp_rank + 1) * shard_size A = A[:, start_idx:end_idx].contiguous() return A - def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + def slice_lora_b_weights(self, B: torch.Tensor): return B @@ -906,10 +921,10 @@ class ReplicatedLinearWithLoRA(BaseLayerWithLoRA): output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None return output, output_bias - def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + def slice_lora_a_weights(self, A: torch.Tensor): return A - def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + def slice_lora_b_weights(self, B: torch.Tensor): return B @@ -1166,10 +1181,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): return final_hidden_states - def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + def slice_lora_a_weights(self, A: torch.Tensor): return A - def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + def slice_lora_b_weights(self, B: torch.Tensor): return B def slice_moe_lora_a_weights( diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 217bd80c3..c462494ad 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -46,6 +46,7 @@ from sglang.srt.lora.utils import ( ) from sglang.srt.managers.io_struct import LoRAUpdateOutput from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import ServerArgs from sglang.srt.utils import get_available_gpu_memory, replace_submodule from sglang.srt.utils.hf_transformers_utils import AutoConfig @@ -82,6 +83,9 @@ class LoRAManager: self.device: torch.device = next(self.base_model.parameters()).device self.tp_size: int = tp_size self.tp_rank: int = tp_rank + # Attention projections shard on the attn-TP group; extracted once + # here (parallel groups are frozen after init_torch_distributed). + self.attn_tp_size: int = get_parallel().attn_tp_size self.lora_added_tokens_size: Optional[int] = None self.enable_lora_overlap_loading: Optional[bool] = ( server_args.enable_lora_overlap_loading @@ -853,6 +857,7 @@ class LoRAManager: dtype=self.dtype, tp_size=self.tp_size, tp_rank=self.tp_rank, + attn_tp_size=self.attn_tp_size, max_lora_rank=self.max_lora_rank, target_modules=self.target_modules, base_model=self.base_model, diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index 587d3a620..3668317c7 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -26,6 +26,7 @@ from sglang.srt.lora.lora import LoRAAdapter from sglang.srt.lora.lora_config import LoRAConfig from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.lora.utils import ( + ATTN_TP_LORA_MODULE_NAMES, EMBEDDING_NAMES, REPLICATED_LINEAR_LORA_NAMES, ROW_PARALLELISM_LINEAR_LORA_NAMES, @@ -137,6 +138,7 @@ class LoRAMemoryPool: dtype: torch.dtype, tp_size: int, tp_rank: int, + attn_tp_size: int, max_lora_rank: int, target_modules: Set[str], base_model: torch.nn.Module, @@ -182,10 +184,15 @@ class LoRAMemoryPool: # here would yield a 4x-narrower inner dim than the adapter weight # (which MoE LoRA modules correctly skip-slice when # `moe_tp_size <= 1`), producing a shape-mismatch - # assert during weight load. Non-MoE modules still shard by - # `tp_size` because attention TP is unchanged. + # assert during weight load. self.moe_tp_size, self.moe_tp_rank = _get_moe_tp_context() + # Attention projections shard along the attention TP group, which + # under `--enable-dp-attention` is `attn_tp_size = tp_size // dp_size`. + # The corresponding LoRA wrappers slice weights by the base layer's + # attn_tp-local rank, so the buffer shapes must match that shard. + self.attn_tp_size: int = attn_tp_size + # Initialize eviction policy self.eviction_policy = get_eviction_policy(eviction_policy) @@ -260,6 +267,20 @@ class LoRAMemoryPool: """Whether this buffer belongs to the shared-expert MoE namespace.""" return module_name.endswith("_shared_moe") + def _effective_tp_size(self, module_name: str) -> int: + """TP width the module's weights are actually sharded along: routed + MoE experts shard by `moe_tp_size` (shared experts by the outer + `tp_size` at EP=1), attention projections by `attn_tp_size` (smaller + than the outer `tp_size` under `--enable-dp-attention`), everything + else by the outer `tp_size`.""" + if self.is_moe_module(module_name) and not self.is_shared_moe_module( + module_name + ): + return self.moe_tp_size + if module_name in ATTN_TP_LORA_MODULE_NAMES: + return self.attn_tp_size + return self.tp_size + @staticmethod def _get_num_experts(base_model: torch.nn.Module) -> int: cfg = base_model.config @@ -388,13 +409,7 @@ class LoRAMemoryPool: module_name, self.base_hf_config, base_model, layer_idx ) c = get_stacked_multiply(module_name, base_model) - # Routed MoE shards along moe_tp_size; shared MoE shards over full TP at EP=1. - effective_tp_size = ( - self.tp_size - if not self.is_moe_module(module_name) - or self.is_shared_moe_module(module_name) - else self.moe_tp_size - ) + effective_tp_size = self._effective_tp_size(module_name) if ( effective_tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES @@ -491,13 +506,8 @@ class LoRAMemoryPool: _, output_dim = get_hidden_dim( module_name, self.base_hf_config, base_model, layer_idx ) - # Same TP-vs-moe-TP sharding rule as get_lora_A_shape above. - effective_tp_size = ( - self.tp_size - if not self.is_moe_module(module_name) - or self.is_shared_moe_module(module_name) - else self.moe_tp_size - ) + # Same sharding rule as get_lora_A_shape above. + effective_tp_size = self._effective_tp_size(module_name) if ( effective_tp_size > 1 and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES @@ -1067,7 +1077,7 @@ class LoRAMemoryPool: # Handle standard modules temp_A_buffer[target_module] = module.slice_lora_a_weights( - temp_A_buffer[target_module], self.tp_rank + temp_A_buffer[target_module] ) cache_keys = temp_A_cache_keys[target_module] assert cache_keys is not None @@ -1077,7 +1087,7 @@ class LoRAMemoryPool: ) temp_B_buffer[target_module] = module.slice_lora_b_weights( - temp_B_buffer[target_module], self.tp_rank + temp_B_buffer[target_module] ) cache_keys = temp_B_cache_keys[target_module] assert cache_keys is not None @@ -1413,7 +1423,7 @@ class LoRAMemoryPool: # Slice B along vocab dimension for this TP rank if self.tp_size > 1: lora_b_weights = lora_lm_head_module.slice_lora_b_weights( - lora_b_weights, self.tp_rank + lora_b_weights ) cache_key = append_cache_key_suffix(name, f"tp{self.tp_rank}") else: diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index 16459ea78..c3a74a515 100644 --- a/python/sglang/srt/lora/utils.py +++ b/python/sglang/srt/lora/utils.py @@ -340,6 +340,24 @@ REPLICATED_LINEAR_LORA_NAMES = [ "fc2_latent_proj", *DSA_INDEXER_LORA_NAMES, ] +# Attention-projection LoRA modules shard on the attention-TP group, which +# under `--enable-dp-attention` is `attn_tp_size = tp_size // dp_size` rather +# than the outer TP size. in_proj / in_proj_qkvz (linear-attention hybrids) +# belong here too: their layers are built on the attn-TP group (mamba.py and +# qwen3_5.py take tp_size/tp_rank from attn_tp when dp attention is enabled). +ATTN_TP_LORA_MODULE_NAMES = frozenset( + { + "qkv_proj", + "qkvr", + "q_b_proj", + "kv_b_proj", + "o_proj", + "out_proj", + "wo_ud", + "in_proj", + "in_proj_qkvz", + } +) # Normalized module names that the LoRA system fully supports # (i.e. get_hidden_dim, init_buffers, and init_lora_modules can handle them). diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index b4fa41ef6..5cca6e1b9 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -123,6 +123,27 @@ _COMMUNICATOR_SPECS = [ ] +def _merge_lora_update_results(results: List[LoRAUpdateOutput]) -> LoRAUpdateOutput: + """Merge the per-rank replies of a LoRA load/unload fan-out into one result. + + The operation succeeded only if every rank succeeded. Reporting a partial + failure as success would let the tokenizer-side LoRA registry drift from + the ranks that failed, so failures win: their deduplicated error messages + are joined, and loaded_adapters reflects the first failed rank. + """ + failed = [r for r in results if not r.success] + if not failed: + return results[0] + error_messages = list( + dict.fromkeys(r.error_message for r in failed if r.error_message) + ) + return LoRAUpdateOutput( + success=False, + error_message=" | ".join(error_messages), + loaded_adapters=failed[0].loaded_adapters, + ) + + class TokenizerControlMixin: """Mixin for TokenizerManager's control-plane operations (weights, cache, lora, profile, internal state, etc.) -- everything that talks to the scheduler via @@ -557,7 +578,9 @@ class TokenizerControlMixin: # Initiate the actual unloading operation at the backend processes only after all # ongoing requests using this LoRA adapter are finished. await self.lora_registry.wait_for_unload(lora_id) - result = (await self.update_lora_adapter_communicator(obj))[0] + result = _merge_lora_update_results( + await self.update_lora_adapter_communicator(obj) + ) return result @@ -574,11 +597,9 @@ class TokenizerControlMixin: "LoRA is not enabled. Please set `--enable-lora` to enable LoRA." ) - # TODO (lifuhuang): Remove this after we verify that dynamic lora loading works - # with dp_size > 1. assert ( - self.server_args.dp_size == 1 - ), "dp_size must be 1 for dynamic lora loading" + self.server_args.dp_size == 1 or self.server_args.enable_dp_attention + ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" logger.info( "Start load Lora adapter. Lora name=%s, path=%s", obj.lora_name, @@ -595,7 +616,9 @@ class TokenizerControlMixin: # Trigger the actual loading operation at the backend processes. obj.lora_id = new_adapter.lora_id - result = (await self.update_lora_adapter_communicator(obj))[0] + result = _merge_lora_update_results( + await self.update_lora_adapter_communicator(obj) + ) # Register the LoRA adapter only after loading is successful. if result.success: @@ -671,7 +694,9 @@ class TokenizerControlMixin: pinned=obj.pinned, ) obj.lora_id = new_adapter.lora_id - result = (await self.update_lora_adapter_communicator(obj))[0] + result = _merge_lora_update_results( + await self.update_lora_adapter_communicator(obj) + ) if result.success: await self.lora_registry.register(new_adapter) @@ -730,11 +755,9 @@ class TokenizerControlMixin: obj.lora_name is not None ), "lora_name must be provided to unload LoRA adapter" - # TODO (lifuhuang): Remove this after we verify that dynamic lora loading works - # with dp_size > 1. assert ( - self.server_args.dp_size == 1 - ), "dp_size must be 1 for dynamic lora loading" + self.server_args.dp_size == 1 or self.server_args.enable_dp_attention + ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" logger.info( "Start unload Lora adapter. Lora name=%s", obj.lora_name, diff --git a/test/registered/unit/lora/test_mem_pool_ep_unit.py b/test/registered/unit/lora/test_mem_pool_ep_unit.py index 280061d34..cb157b13f 100644 --- a/test/registered/unit/lora/test_mem_pool_ep_unit.py +++ b/test/registered/unit/lora/test_mem_pool_ep_unit.py @@ -101,10 +101,10 @@ class _FakeRoutedMoeLayer(_FakeFusedMoEWithLoRA, _IdentityMoeSlices): class _FakeDenseLayer: - def slice_lora_a_weights(self, weights, _rank): + def slice_lora_a_weights(self, weights): return weights - def slice_lora_b_weights(self, weights, _rank): + def slice_lora_b_weights(self, weights): return weights @@ -994,6 +994,7 @@ class TestPoolInitPicksUpEpContext(unittest.TestCase): dtype=torch.bfloat16, tp_size=tp_size, tp_rank=tp_rank, + attn_tp_size=tp_size, max_lora_rank=8, target_modules={"qkv_proj"}, base_model=base_model, @@ -1089,6 +1090,9 @@ def _fake_base_model_with_hidden_dim(num_experts: int) -> torch.nn.Module: return cfg.hidden_size, cfg.moe_intermediate_size * 2 if module_name == "down_proj_moe": return cfg.moe_intermediate_size, cfg.hidden_size + if module_name == "in_proj_qkvz": + # linear-attention qkvz input projection (column-parallel) + return cfg.hidden_size, 4 * cfg.hidden_size raise NotImplementedError(module_name) return _Model() @@ -1119,6 +1123,9 @@ class TestMoeBufferShardsByMoeTp(unittest.TestCase): pool.max_loras_per_batch = 2 pool.tp_size = tp_size pool.tp_rank = 0 + # Without --enable-dp-attention the attention TP group equals the + # outer TP group. + pool.attn_tp_size = tp_size pool.moe_ep_size = ep_size pool.moe_ep_rank = ep_rank pool.moe_tp_size = moe_tp_size @@ -1218,6 +1225,76 @@ class TestMoeBufferShardsByMoeTp(unittest.TestCase): self.assertEqual(q_b, (2, 48, 8)) +class TestAttnModulesShardByAttnTp(unittest.TestCase): + """Regression: attention-module LoRA buffers must shard by `attn_tp_size`, + not the outer `tp_size`. + + Under `--enable-dp-attention` attention layers are built on the attn_tp + group (`attn_tp_size = tp_size // dp_size`), so e.g. MLA `o_proj` holds an + attn_tp-local input shard. Sizing the LoRA buffer by the outer `tp_size` + would make it narrower than the slice produced by + `RowParallelLinearWithLoRA.slice_lora_a_weights` (which slices by the base + layer's attn_tp-local rank), failing the shape-match assert at load time. + """ + + def _pool(self, *, tp_size: int, attn_tp_size: int) -> LoRAMemoryPool: + pool = LoRAMemoryPool.__new__(LoRAMemoryPool) + pool.max_loras_per_batch = 2 + pool.tp_size = tp_size + pool.tp_rank = 0 + pool.attn_tp_size = attn_tp_size + pool.moe_ep_size = 1 + pool.moe_ep_rank = 0 + pool.moe_tp_size = tp_size + pool.moe_tp_rank = 0 + pool.moe_use_local_expert_ids = False + pool._num_experts_local = 1 + pool.experts_shared_outer_loras = False + pool.base_hf_config = types.SimpleNamespace( + hidden_size=64, + num_attention_heads=8, + num_key_value_heads=8, + head_dim=8, + intermediate_size=256, + moe_intermediate_size=192, + ) + return pool + + def test_attn_tp_1_keeps_attention_buffers_full_width(self): + """tp=4 with attn_tp=1 (dp-attention, dp=4): attention weights are + replicated across ranks, so the LoRA buffers must be full-width. + """ + pool = self._pool(tp_size=4, attn_tp_size=1) + model = _fake_base_model_with_hidden_dim(num_experts=1) + # o_proj is row-parallel: A input_dim = head_dim*num_heads = 64, + # undivided under attn_tp=1 (pre-fix: 16). + self.assertEqual(pool.get_lora_A_shape("o_proj", model, 8, 0), (2, 8, 64)) + # qkv_proj is column-parallel: B output_dim = 8 * 24 = 192, + # undivided under attn_tp=1 (pre-fix: 48). + self.assertEqual(pool.get_lora_B_shape("qkv_proj", model, 8, 0), (2, 192, 8)) + + def test_attn_tp_gt1_still_shards_attention_buffers(self): + """tp=4 with attn_tp=2 (dp=2): attention weights are sharded 2-way.""" + pool = self._pool(tp_size=4, attn_tp_size=2) + model = _fake_base_model_with_hidden_dim(num_experts=1) + self.assertEqual(pool.get_lora_A_shape("o_proj", model, 8, 0), (2, 8, 32)) + self.assertEqual(pool.get_lora_B_shape("qkv_proj", model, 8, 0), (2, 96, 8)) + + def test_linear_attention_in_proj_shards_by_attn_tp(self): + """Regression: in_proj_qkvz is built on the attn-TP group under + dp-attention (qwen3_5.py passes tp_rank=attn_tp_rank), but it was + classified as outer-TP, so with tp=4 / attn_tp=1 its LoRA B buffer + came out 4x narrower than the wrapper's attn_tp-local slice and + adapter load failed on the shape assert.""" + pool = self._pool(tp_size=4, attn_tp_size=1) + model = _fake_base_model_with_hidden_dim(num_experts=1) + # column-parallel: B output_dim = 4*64 = 256, undivided under + # attn_tp=1 (pre-fix: divided by the outer tp=4 -> 64). + self.assertEqual( + pool.get_lora_B_shape("in_proj_qkvz", model, 8, 0), (2, 256, 8) + ) + + class TestLoadBufferPassesMoeTpRankToSlice(unittest.TestCase): """Regression: `load_lora_weight_to_buffer` must hand `moe_tp_rank` (not the outer `tp_rank`) to `slice_moe_lora_{a,b}_weights`. diff --git a/test/registered/unit/managers/test_lora_update_result_merge.py b/test/registered/unit/managers/test_lora_update_result_merge.py new file mode 100644 index 000000000..fb4cc221b --- /dev/null +++ b/test/registered/unit/managers/test_lora_update_result_merge.py @@ -0,0 +1,78 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Unit tests for merging per-rank LoRA update replies from the control fan-out.""" + +import unittest + +from sglang.srt.managers.io_struct import LoRAUpdateOutput +from sglang.srt.managers.tokenizer_control_mixin import _merge_lora_update_results +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +def _ok(adapters=None) -> LoRAUpdateOutput: + return LoRAUpdateOutput(success=True, loaded_adapters=adapters or {}) + + +def _err(message, adapters=None) -> LoRAUpdateOutput: + return LoRAUpdateOutput( + success=False, error_message=message, loaded_adapters=adapters or {} + ) + + +class TestMergeLoRAUpdateResults(CustomTestCase): + def test_all_success_returns_first_rank_result(self): + """On success the merge must hand back a rank's own reply: callers + mutate result.loaded_adapters in place during LRU eviction, which a + synthesized empty result would silently break.""" + results = [_ok({"a": "path"}), _ok({"a": "path"})] + merged = _merge_lora_update_results(results) + self.assertIs(merged, results[0]) + self.assertTrue(merged.success) + + def test_any_rank_failure_wins(self): + """Regression guard for the pre-merge behavior of returning + results[0]: a failure on a non-zero rank was reported as success, + letting the tokenizer-side registry drift from that rank's actual + adapter state.""" + merged = _merge_lora_update_results( + [_ok({"a": "path"}), _err("out of memory", {"stale": "path"})] + ) + self.assertFalse(merged.success) + self.assertEqual(merged.error_message, "out of memory") + self.assertEqual(merged.loaded_adapters, {"stale": "path"}) + + def test_duplicate_error_messages_deduplicated(self): + """All ranks usually fail identically (e.g. "already loaded"); the + joined message must not repeat per rank, but distinct causes must all + be kept.""" + merged = _merge_lora_update_results( + [_err("already loaded"), _err("already loaded"), _err("bad rank")] + ) + self.assertFalse(merged.success) + self.assertEqual(merged.error_message, "already loaded | bad rank") + + def test_failure_without_message(self): + """A rank replying success=False with error_message=None must not + crash the join.""" + merged = _merge_lora_update_results([_err(None), _ok()]) + self.assertFalse(merged.success) + self.assertEqual(merged.error_message, "") + + +if __name__ == "__main__": + unittest.main()