[peft] fix: fall back for unaligned dense multi-LoRA grouped MM

[yushengsu-thu]

Summary

• gate each dense MultiLoRALinear projection on PyTorch grouped-MM's 16-byte address, stride, dtype, and output-layout requirements
• fall back to per-slot F.linear only for unsafe A or B projections while preserving the existing TP/SP collective placement and the aligned grouped-MM fast path
• retain host token counts at the existing dispatch-time synchronization so fallback splits follow sequence-parallel windows and zero-token slots remain in the autograd graph

Root cause

With TP8 and a physical max LoRA rank of 16, the column-parallel A projection has a local BF16 rank of 2. The TP gather's backward produces a contiguous [tokens, 2] gradient with a 4-byte leading stride, while PyTorch grouped-MM requires non-unit strides and operand addresses to be 16-byte aligned. The failure therefore occurs inside grouped-MM backward before the optimizer runs:

RuntimeError: strides should be multiple of 16 bytes

The selector uses the physical local weight shape and current tensor layout, not the adapter's logical rank. Unsafe inputs use the mathematically equivalent per-slot linear path for both projections. Every slot is still evaluated, including empty slots, so Megatron gradient hooks observe present zero gradients.

Validation

Validated on NVIDIA H200 with PyTorch 2.11.0+cu130:

• exact TP8 / physical max-rank 16 old-vs-new reproduction:
  • origin/bridge reproduced the 16-byte stride exception
  • this branch passed output, input-gradient, A-gradient, and B-gradient parity against a per-slot reference
  • dispatch counters were fallback=1, grouped_mm=1, proving the unaligned local A uses fallback while the gathered, aligned B keeps the fast path
• TP8 / physical max-rank 64 passed parity with fallback=0, grouped_mm=2
• TP8 / max-rank 16 with activation recompute passed parity with the expected replay counts
• 84 passed across:
  • tests/unit_tests/peft/test_multi_lora_layers.py
  • tests/unit_tests/peft/test_multi_lora.py
  • tests/unit_tests/peft/test_multi_lora_moe.py
• focused regressions cover physical BF16 rank 2, misaligned storage offsets, FP32 fallback, an all-zero batch, sequence-parallel window narrowing across a slot boundary, aligned fast-path parity, and zero-token slot gradients
• uv run --no-project pre-commit run --all-files passed

Scope

This PR fixes the dense MultiLoRALinear path. MultiLoRAGroupedExpertLinear is unchanged; its physical rank alignment remains a separate concern. The validation covers the exact TP8 dense layer and collective path, but it is not a fresh full 744B GLM-5.2 end-to-end rerun at max-rank 16.

PyTorch grouped-MM references:

• https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/GroupedMMUtils.h
• https://github.com/pytorch/pytorch/blob/main/torch/csrc/autograd/FunctionsManual.cpp#L1760-L1812

[yushengsu-thu]

/ok to test e93ef87

[yushengsu-thu]

Follow-up 9872e0f applies the agreed polish: the alignment constant is now _PYTORCH_GROUPED_MM_ALIGNMENT_BYTES, grouped with _GROUPED_MM_SUPPORTED_DTYPES under one contract comment citing PyTorch ATen's GroupedMMUtils.h (https://github.com/pytorch/pytorch/blob/ab5fb26f8ffc6e4dc97b51b5611bce957645b1db/aten/src/ATen/native/GroupedMMUtils.h#L24-L48) — it restates that build's layout contract, not a tunable, and the comment records the failure asymmetry (too strict -> safe fallback; too loose -> loud host-side TORCH_CHECK). Tests now import the constant instead of a second % 16 literal, and two tests were added: a GPU boundary canary that builds a BF16 input exactly 16- but not 32-byte aligned (offset 8 elements from the allocator's 256B base — allocator-natural tensors could never probe this) and asserts _can_use_grouped_mm admits it and torch._grouped_mm matches the per-slot F.linear reference, so any upstream tightening of the contract fails loudly here; plus a seeded 50-case equivalence check for the device/host _narrow_token_counts_to_window twins. Validated on an H200 (torch 2.11.0+cu130): full file 51 passed with GPU, 39 passed / 12 skipped CPU-only; ruff check/format clean.