fix(peft): support mixed-dtype memory-efficient LoRA backward (#3675)

* fix(peft): support mixed-dtype memory-efficient LoRA backward

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* fix(peft): support mixed-dtype forward/backward in the fused LoRA MLPs

`apply_lora_to_linear_modules()` installs the fused MLP wrapper whenever all of
gate/up/down are LoRA-patched, so production execution enters
`LoRASwiGLUMLPFunction` / `LoRAReLU2MLPFunction` rather than the per-linear
`LoRATritonFunction`. Both fused functions rejected mixed-dtype layouts that the
per-linear path handles:

* backward runs outside autocast, so with FP32 activations and BF16 compute
  weights (FSDP2 `param_dtype=bf16` + `output_dtype=fp32`) the recomputed
  matmuls raised "expected m1 and m2 to have the same dtype", and the input
  gradient came back in the compute dtype instead of the input's;
* forward builds each projection with `addmm_`, which is not autocast-eligible
  (an in-place op cannot change its output dtype), so a `lora_dtype` differing
  from the base weights raised "self and mat2 must have the same dtype".

Backward now recomputes in the forward compute dtype (the saved `e` is a forward
matmul output, so its dtype is exactly what autocast picked) and returns every
gradient in its input's dtype, matching what autograd does through autocast's
own cast nodes; forward casts the `addmm_` operands to the base projection's
dtype. All casts are no-ops when the dtypes already agree, so the uniform-dtype
paths stay bitwise unchanged and allocate nothing extra.

Add a regression test that goes through `apply_lora_to_linear_modules()` for
both a SwiGLU and a ReLU² MLP, over both mixed-dtype layouts, asserting the
fused autograd node is really the one exercised and checking parity plus
gradient dtypes against the per-linear path.

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* test(peft): cover mixed-dtype LoRA backward on two-rank FSDP2

The unit tests construct the FP32-activation/BF16-weight layout by hand. What
actually broke in the field is that a real `MixedPrecisionPolicy` *produces* it:
FSDP2's `output_dtype=float32` cast is what puts an FP32 activation in front of
BF16 weights, and it only exists inside a sharded, multi-rank graph.

Add a two-rank functional test under the CI-collected `hf_peft` folder that
shards a LoRA-patched block with the exact policy from issue #3652
(`param_dtype=bf16`, `reduce_dtype=fp32`, `output_dtype=fp32`) and covers all
three memory-efficient paths: the per-linear `LoRATritonFunction` from the
issue's traceback, and the `LoRASwiGLUMLPFunction` / `LoRAReLU2MLPFunction` that
`apply_lora_to_linear_modules()` installs automatically. Each case asserts the
mixed dtype was actually observed, so the test cannot decay into uniform-dtype
coverage, then compares outputs and gradients against the same model on plain
autograd under the identical sharded setup.

Verified on two GPUs: fails on `main` with the issue's own
"self and mat2 must have the same dtype, but got Float and BFloat16", fails on
the per-linear-only fix with "expected m1 and m2 to have the same dtype", and
passes with both.

Also document the tensor contracts of the two new `lora_mlp` helpers.

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* fix(peft): decline the Triton LoRA kernels on mixed dtype

The Triton path is the third route into memory-efficient LoRA and had the same
defect as the two already fixed here. `lora_forward_wrapper` feeds its operands
straight to `tl.dot`, which asserts a single dtype, and the kernels run outside
autocast, so an FP32 activation meeting BF16 adapters aborts compilation in the
*forward*:

    AssertionError: Both operands must be same dtype. Got fp32 and bf16
    triton.compiler.errors.CompilationError

This is the layout from issue #3652, whose recipe does set `use_triton: true` —
it only escaped the Triton kernels because `patch_linear_module` force-disables
them when TransformerEngine is importable. Without TE (a supported install), 91
example recipes that set `use_triton: true` hit this.

Decline the kernels when their own precondition does not hold and fall through
to the torch matmul path, which follows the forward compute dtype and casts
explicitly. Verified on two GPUs: the Triton kernels are still used for uniform
dtypes (no perf change), and in the mixed case the result is bit-identical to
`use_triton: false` — output and all gradients.

Cover the routing decision with a CPU unit test, and add a `per_linear_triton`
case to the two-rank FSDP2 functional test, skipped when TE is present since
Triton is unreachable then.

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* fix(peft): decline fused LoRA MLP for biased projections

The fused MLP forward builds each projection with `F.linear(x, base_weight)` and
never passes a bias, so a biased base projection is trained on silently wrong
math. Models plumb this straight from the HF config
(`nn.Linear(..., bias=config.mlp_bias)` in llama/qwen2/nemotron_v3), and
`moe.layers.MLP` takes `bias` as a first-class argument.

The bug predates this PR — on `main` a `mlp_bias: true` SwiGLU MLP already
diverges from the per-linear path by 1.51 in plain FP32. But it was masked in
the mixed-precision layout this PR targets, where the backward used to raise
before the wrong numbers could matter; fixing the dtype handling turns that
crash into a silent divergence. Gate it in `_fusible`, which is already where
fusion is declined for DoRA, active dropout, quantized bases, DTensor, and meta
weights. Biased MLPs fall back to the per-linear path and match it exactly
(delta 1.51 -> 0.0); unbiased MLPs still fuse.

Also cast `gW`/`uW` inside the `needs_x` branch that is their only reader —
eagerly casting two `(inter, hidden)` base weights costs real memory on every
backward that does not need `d_x` — and correct `_cast`'s docstring, which
invited callers to mutate a result that may alias the caller's `grad_out`.

Uniform-dtype paths remain bitwise identical to pre-PR.

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* test(peft): close blind spots in the mixed-dtype LoRA tests

An adversarial pass over the new tests found several assertions that could not
fail and one case that could never run. Each item below was confirmed by
injecting the corresponding bug and watching the suite stay green.

* `alpha == dim` made `PeftConfig.scale` exactly 1.0, so every dropped-scale bug
  was invisible: deleting `dS` from the fused backward passed both the new
  functional and unit tests. Use `alpha = 3 * dim` (the shipped default is
  `dim=8, alpha=32`, and ~39 example recipes use `alpha != dim`). The same
  injection now fails the functional test and 6 unit tests.

* The `per_linear_triton` case was dead in CI. `patch_linear_module` disables
  Triton whenever TransformerEngine is importable and the CI image builds it
  (`docker/Dockerfile` `TE_COMMIT`), so the case was skipped by a bare `print` --
  no skip marker, nothing that could fail -- and the entire Triton fix could be
  deleted with CI still green. Drop the case and cover the fallback in
  `test_lora_kernel.py` by calling `apply_memory_efficient_lora` directly, which
  runs regardless of TE and asserts bit-exact agreement with the torch path.
  Removing the guard now fails that test.

* The gradient-dtype assertions could not fail: torch's autograd engine coerces
  a Function's returned gradient to its input's dtype (verified on 2.10). Drop
  them and record why, so they do not come back as false confidence.

* Nothing required the baseline to be a *different* path, so a refactor making
  fusion unconditional would have left every parity check comparing an
  implementation against itself -- and passing. Observe the baseline too and
  require it entered no memory-efficient autograd function.

* Tolerance is now `3e-2` with the measurement behind it recorded: the fused vs
  per-linear floor is bf16 GEMM ordering, <=8.8e-3 on sm_89, and CI also runs
  A100/H100/GB200. Detection is unaffected (the bugs this catches land at
  >=5.0e-2).

* Stop forcing `CUDA_VISIBLE_DEVICES=0,1` in the worker env. The skip gate counts
  the *parent's* visible devices, so on a host where the scheduler assigned other
  GPUs this would have pushed the ranks onto another job's. `run_test.sh` already
  exports the right value.

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

---------

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>