@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 32b7177909

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 13:00 UTC

[LoRA][Bugfix] Dedup LoRA wrapping for modules referenced from multiple attribute paths (MoE gate) (#42757)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed May 16, 2026 at 11:22 UTC 32b7177909d1c9928bcedd81de7de5a1fa21d2b3

2 files changed +115

tests/lora/test_lora_manager.py

+94

@@ -179,6 +179,100 @@ def test_wrap_gate_linear(default_vllm_config, dist_init, dummy_model):

179 )

180

181

182 +def test_dedup_shared_module_across_paths(default_vllm_config, dist_init, dummy_model):

183 + """A module reachable from two attribute paths (e.g. a MoE gate held

184 + both directly on the block and inside its inner runner) must produce a

185 + single LoRA wrapper. Both paths must end up pointing to that same

186 + wrapper instance, and only the canonical path should live in

187 + `manager.modules` — otherwise activate_adapter would call `reset_lora`

188 + on the alias and clobber weights set under the canonical name.

189 + """

190 + from vllm.model_executor.layers.linear import ReplicatedLinear

191 +

192 + class AliasContainer(nn.Module):

193 + def __init__(self, gate: nn.Module):

194 + super().__init__()

195 + self.gate = gate # canonical path: "moe.gate"

196 +

197 + # Inner submodule holding the SAME gate instance under another

198 + # path. This mirrors how FusedMoE.runner.gate references the

199 + # block's gate in qwen3_moe.

200 + class _Runner(nn.Module):

201 + def __init__(self, g):

202 + super().__init__()

203 + self.gate = g # alias path: "moe.runner.gate"

204 +

205 + self.runner = _Runner(gate)

206 +

207 + gate = ReplicatedLinear(10, 4, bias=False)

208 + model = dummy_model

209 + model.add_module("moe", AliasContainer(gate))

210 +

211 + assert model.moe.gate is model.moe.runner.gate

212 +

213 + manager = LoRAModelManager(

214 + model,

215 + 1,

216 + 1,

217 + 1,

218 + LoRAConfig(

219 + max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE

220 + ),

221 + torch.device(DEVICES[0]),

222 + )

223 +

224 + canonical = manager.model.get_submodule("moe.gate")

225 + alias = manager.model.get_submodule("moe.runner.gate")

226 +

227 + # Same wrapper instance on both paths so forward through either side

228 + # sees the LoRA-augmented module.

229 + assert isinstance(canonical, ReplicatedLinearWithLoRA)

230 + assert alias is canonical

231 +

232 + # Only the canonical path is tracked as a LoRA target. Tracking the

233 + # alias would cause activate_adapter to reset_lora on it after the

234 + # canonical entry already populated the weights.

235 + assert "moe.gate" in manager.modules

236 + assert "moe.runner.gate" not in manager.modules

237 +

238 +

239 +def test_lm_head_exempt_from_dedup(default_vllm_config, dist_init, dummy_model):

240 + """The dedup logic must NOT collapse `lm_head` even when it is reachable

241 + from another attribute path (tied-embedding models do

242 + `self.lm_head = self.model.embed_tokens`, sharing the same nn.Module

243 + instance). The lm_head branch additionally rewires `logits_processor`

244 + into a `LogitsProcessorWithLoRA`, so skipping it would silently break

245 + LoRA on lm_head.

246 + """

247 + from vllm.lora.layers import LogitsProcessorWithLoRA

248 +

249 + # Add a non-lm_head alias to the same module instance as lm_head. The

250 + # dedup keys on id(module); without the lm_head exemption the alias

251 + # would consume the wrapped_by_id slot first and lm_head would be

252 + # silently skipped, so logits_processor would never be wrapped.

253 + model = dummy_model

254 + model.add_module("embed_tokens", model.lm_head)

255 + assert model.embed_tokens is model.lm_head

256 +

257 + manager = LoRAModelManager(

258 + model,

259 + 1,

260 + 1,

261 + 1,

262 + LoRAConfig(

263 + max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE

264 + ),

265 + torch.device(DEVICES[0]),

266 + )

267 +

268 + # lm_head's special handling still ran: logits_processor got wrapped

269 + # and the lm_head entry is tracked under self.modules.

270 + assert isinstance(

271 + manager.model.get_submodule("logits_processor"), LogitsProcessorWithLoRA

272 + )

273 + assert "lm_head" in manager.modules

274 +

275 +

276 def test_skip_unsupported_matched_modules(default_vllm_config, dist_init, dummy_model):

277 class UnsupportedContainer(nn.Module):

278 def __init__(self):

vllm/lora/model_manager.py

+21

@@ -361,6 +361,8 @@ class LoRAModelManager:

361 # - given an input 'x' return ''

362 return module_name.rpartition(".")[0]

363

364 + wrapped_by_id: dict[int, BaseLayerWithLoRA] = {}

365 +

366 for module_name, module in self.model.named_modules(remove_duplicate=False):

367 if isinstance(module, PPMissingLayer):

368 continue

@@ -391,6 +393,23 @@ class LoRAModelManager:

393 )

394 continue

395

396 + existing_wrapper = wrapped_by_id.get(id(module))

397 + if existing_wrapper is not None and "lm_head" not in module_name:

398 + # Same underlying module was already wrapped under another

399 + # path (e.g. a MoE gate held both directly on the block and

400 + # inside the MoE runner). The adapter targets the canonical

401 + # path (`mlp.gate`); rewire the alias attribute

402 + # (`runner.gate`) to the SAME wrapper so the forward path

403 + # through the alias still applies LoRA, but do NOT add a

404 + # second entry to self.modules — otherwise `activate_adapter`

405 + # would call `reset_lora` on the alias and wipe the weights

406 + # just set under the canonical name, because the alias can't

407 + # load LoRA weights due to name mismatch.

408 + parent = self.model.get_submodule(_parent_module(module_name))

409 + # reference

410 + setattr(parent, module_name.rpartition(".")[-1], existing_wrapper)

411 + continue

412 +

413 parts = module_name.split(".")[-1]

414 packed_moduled_lst = self.packed_modules_mapping.get(parts, [])

415 if isinstance(module, FusedMoE):

@@ -411,6 +430,8 @@ class LoRAModelManager:

430 self.model.config,

431 ),

432 )

433 + if isinstance(new_module, BaseLayerWithLoRA):

434 + wrapped_by_id[id(module)] = new_module

435

436 # (yard1): TODO make this more robust

437 if "lm_head" in module_name: