[peft] fix: Guard dense Multi-LoRA grouped MM layouts

[yaoyu-33]

What does this PR do?

Guards dense Multi-LoRA's private grouped-MM calls with the installed kernel's forward/backward contract. Eligible CUDA layouts retain the grouped-MM fast path; unsupported layouts, architectures, dtypes, empty batches, and any zero-count adapter slot use per-slot torch.nn.functional.linear.

This is a narrow correctness follow-up to the supported Multi-LoRA path introduced in #4218. A downstream production report reproduced the failure at TP8 with a BF16 local rank of 2: radixark/Megatron-Bridge#32. The governing 16-byte pointer and stride requirements are defined by PyTorch GroupedMMUtils.h.

Supported trigger and impact

Dense Multi-LoRA may call the private grouped-MM operation for both A and B projections. TP can shard the physical BF16 rank to a width whose contiguous backward gradient has a row stride below the required 16-byte boundary. Misaligned views, FP32 autograd, pre-SM80 devices, empty inputs, and zero-count slots are also ineligible. Passing a mixed layout such as [2, 0, 2] to grouped-MM can make backward unsafe before the first optimizer update.

Root cause

The prior path assumed that valid logical LoRA shapes implied valid grouped-MM storage and routing layouts. PyTorch additionally requires aligned pointers, aligned batch and matrix strides, supported autograd dtypes, non-empty groups, an aligned output-gradient row stride, and an SM80-or-newer CUDA device.

Changelog

• Check grouped-MM pointer, stride, dtype, device, architecture, empty-input, zero-split, and backward-output-width requirements before each dense A/B projection.
• Force the dense per-slot fallback whenever any adapter count is zero, preserving inactive weights' explicit zero-gradient dependencies.
• Cache immutable host split sizes once in set_tokens_per_adapter_slot(), derive SP-local splits on the host, and remove per-layer .tolist() synchronization.
• Build grouped offsets once per forward and reuse the same tensor for A and B.
• Preserve TP/SP collective placement, scaling, slot ordering, and the eligible grouped-MM fast path.
• Add deterministic CPU contracts for mixed zero-slot backward, host split/offset reuse, and SM80 eligibility, plus a real CUDA backward test guarded by actual device availability.

Validation

Fail-before on dd150c1d73ed9a25ad3ee87149cb98a0a2ca3614 with production code unchanged:

uv run python -m pytest tests/unit_tests/peft/test_multi_lora_layers.py::TestMultiLoRALinearSlots::test_forward_falls_back_when_grouped_mm_layout_is_ineligible -q
1 failed: RuntimeError: strides should be multiple of 16 bytes

Pass-after with the identical command:

1 passed

Correction-focused contracts:

4 passed, 1 skipped

The skipped test is the real SM80+ CUDA mixed-zero-slot backward test. The available host driver reports CUDA 12.4 and is too old for the installed PyTorch CUDA runtime, so CUDA could not initialize; no GPU result is simulated.

Focused adjacent coverage:

uv run python -m pytest tests/unit_tests/peft/test_multi_lora.py tests/unit_tests/peft/test_multi_lora_layers.py -q
60 passed, 6 skipped

Additional gates:

git diff --check
passed

uv run pre-commit run --all-files
passed

Non-goals

• No adapter scheduling or lifecycle changes.
• No TP/SP collective or checkpoint-format changes.
• No grouped expert LoRA changes.
• No new kernel policy, dependency, configuration, or public API.

Before your PR is Ready for review

• Read and followed the contributor guidelines.
• Added focused regression tests.
• Documentation is not needed for this internal correctness guard.
• No optional dependency surface is changed.

Additional Information

• Related to #1986.
• Commit: d67579f477ad99c2ed9040eb57d9a77d408b2bf0.

[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.

[yaoyu-33]

/ok to test 16f0e32

[yaoyu-33]

/ok to test d67579f