fix(peft): correct dense MultiLoRALinear on replicated base linears

[yushengsu-thu]

What does this PR do ?

Fixes the dense MultiLoRALinear (#4218) on replicated base linears — TELinear parallel_mode='duplicated', i.e. the MLA q/kv down-projections of GLM-/DeepSeek-style models.

Stacked on #4218 and #5155 — this branch includes their commits as its base, so the Files tab shows all three until they merge. Only the last commit is new here. Ported from radixark#25, where it runs in production on the bridge branch.

Why

A GLM-5.2 (5-layer) multi-LoRA e2e smoke — the first MLA model through the multi-LoRA path — crashed in the first training forward. Qwen3-style models never reach these paths (all their LoRA targets are TP-parallel linears), which is why the existing e2e coverage missed them.

Two defects, both specific to replicated bases

1. Per-slot token spans covered the full micro-batch, but a replicated base under sequence parallelism consumes the SP shard directly (TE stores parallel_mode=None for duplicated linears, so the adapter attribute derivation disables the SP gather). The grouped GEMM's offsets then exceed the row count — an out-of-bounds read surfacing as CUDA error: device-side assert triggered, not a shape error. The spans are now intersected with this rank's contiguous token window (_narrow_token_counts_to_window, the same invariant the MoE slot routing uses for its SP narrow). The host-side total is cached once per micro-batch in set_tokens_per_adapter_slot, so the forward needs no device sync.

2. _gather_output keyed only on input_is_parallel, leaving the adapter output as a [tokens, out/tp] shard while a replicated base produces the full [tokens, out]. It now mirrors ParallelLinearAdapter's lin_out_gather_output: gather for row-parallel and replicated bases.

Deliberately unchanged: the TP collective between the two grouped GEMMs stays unconditional. ParallelLinearAdapter shards A on the rank axis whenever input_is_parallel=False — replicated bases included — so the mid all-gather is required there too (gating it on base_linear_is_parallel was tried and breaks GEMM2 with a contraction mismatch).

Testing

• 4 new CPU tests for the window-narrowing math (slot-boundary-spanning windows, identity, single-slot, many-small-slots); the three multi-LoRA test files pass 74/74 on an H200 devbox (radixark fork, same file contents as this PR).
• GLM-5.2 (5-layer) multi-LoRA e2e (TP=EP=4, ETP=1, SP, thd): with this fix the run completes 3 training steps end-to-end with per-step distributed adapter upserts and (tp, pp, ep) checkpoint shards, executed under CUDA_LAUNCH_BLOCKING=1 so no async assert hides.
• Without the fix, the identical run dies in the first wrapped MLA down-projection (torch._grouped_mm device-side assert).

Changelog

• MultiLoRALinear: narrow per-slot token spans to this rank's SP shard for replicated bases (_narrow_token_counts_to_window); cache the host-side span total in set_tokens_per_adapter_slot
• MultiLoRALinear: gather the column-sharded adapter output for replicated bases (mirror ParallelLinearAdapter.lin_out_gather_output)
• 4 new CPU unit tests for the window narrowing

GitHub Actions CI

See the CI section in the Contributing doc. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

Additional Information

• Depends on [peft] Support multi-lora for Megatron models #4218 and feat(peft): support multi-LoRA on grouped MoE expert linears #5155

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[yaoyu-33]

Hi @yushengsu-thu — this needs a rebase too, and it is stacked on #5155, so please rebase that one first.

Why it conflicts: #4218 was squash-merged into main as 1f12931e2. This branch still carries #4218's six original (unsquashed) commits plus #5155's commit as its base, so git conflicts with the squashed version on main.

Fix: after rebasing #5155 (yusheng/multilora-moe-experts) per my comment there, restack this branch on it. I verified locally that this cherry-picks cleanly on top of the rebased #5155:

git checkout yusheng/multilora-replicated-base
git reset --hard yusheng/multilora-moe-experts   # already rebased onto upstream/main
git cherry-pick c78e227f795bdcb4f66ef14a3489b4f3028edd1b
git push --force-with-lease

FYI this commit does not apply directly onto main — it genuinely depends on the MultiLoRAGroupedExpertLinear code from #5155, so the stack has to be preserved.

I would have pushed the rebase for you, but the fork is owned by the radixark org and GitHub's "allow edits by maintainers" does not apply to org-owned forks, so my push is rejected.