[peft] Support multi-lora for Megatron models

[mathewjhan]

What does this PR do ?

Adds a simple multi-lora implementation with a fixed amount of adapter slots + adds some helper methods to manage the slot lifecycle

Changelog

• Adds MultiLoRA transform class based on PEFT that replaces (only supports non-experts for now)
• Adds MultiLoRALinear which reuse ParallelLinearAdapter in a module list + performs a group gemm in the forward pass
• Reuses current export methods for LoRA adapters by exposing a single ParallelLinear adapter as .adapter
• Adds lifecycle methods for handling adapters in different slots

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. 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?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

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

Review notes from the Multi-LoRA pass. I did not run tests; this is a static review focused on correctness/integration/docs.

1. Adapter construction does not preserve the wrapped Megatron runtime config.

MultiLoRALinear constructs each ParallelLinearAdapter without forwarding the wrapped module's model-parallel config and several parallel-layout flags. Existing LoRA.transform() passes things like model_parallel_config=module.config, disable_tensor_parallel_comm, and base_linear_is_parallel when constructing ParallelLinearAdapter; this path only forwards a subset of those fields.

Impact: real TP/SP, dtype, shared/expert, and initialization behavior can diverge from the wrapped layer. In particular, BF16/FP16 adapter casting and communication behavior can be wrong in real training even if the CPU fake tests pass.

Suggested fix: mirror the adapter kwargs used by the existing LoRA path unless there is a deliberate Multi-LoRA-specific reason not to.

2. Slot rank/alpha metadata is not device-safe, checkpointed, or export-correct.

alpha_values and rank_values are plain tensors, not registered buffers or parameters. They will not follow .cuda()/.to(), and they are not included in state_dict() / sharded_state_dict(). Also, expose_adapter_slot() exposes the raw ParallelLinearAdapter, whose .alpha / .dim remain constructor defaults, so export/merge can use max-rank/default scaling instead of the per-slot rank / alpha initialized through init_adapter_slot().

Impact: a normal model move can leave slot metadata on CPU, checkpoint resume loses slot lifecycle state, and HF export/merge can use the wrong adapter scaling/rank for initialized slots.

Suggested fix: register rank/alpha as persistent buffers, include them in adapter filtering/checkpointing, and make the exposed slot metadata reflect the selected slot. If export always uses max-rank padded tensors, please document that contract explicitly and test it.

3. Distributed checkpoint sharding metadata is bypassed.

MultiLoRALinear.sharded_state_dict() currently inserts self.adapters.state_dict() raw tensors. Existing adapter wrappers delegate to each adapter's sharded_state_dict() so ParallelLinearAdapter can emit the TP/expert/SwiGLU sharding metadata.

Impact: native distributed checkpoint save/load can lose sharding metadata or fail under TP/PP and any future expert/shared-adapter support.

Suggested fix: iterate adapter slots and call each slot's sharded_state_dict(prefix=f"{prefix}adapters.{i}.", ...), plus save the slot metadata buffers.

4. Active slot routing has no usable default or validation.

forward() immediately calls tokens_per_adapter.cumsum(...), but tokens_per_adapter defaults to None, and set_tokens_per_adapter_slot() stores the tensor as-is. There is no validation for length, nonnegative counts, device, dtype, or sum matching the flattened token count.

Impact: a normal forward without a prior routing call crashes with a cryptic AttributeError, and malformed routing can produce grouped-GEMM failures or incorrect slot assignment.

Suggested fix: validate in the setter or forward, move routing metadata to the input device as needed, and either default all tokens to slot 0 or raise a clear ValueError explaining that routing must be set before forward.

5. Lifecycle context managers can leave the model mutated after exceptions.

expose_adapter_slot() and hide_adapters() pop modules from _modules, then restore only after yield; neither restoration is protected by try/finally. An invalid slot index can also raise after adapters is already popped.

Impact: failed export/load/checkpoint code can leave modules missing .adapters or stuck with .adapter, corrupting later training/export in the same process.

Suggested fix: validate slot index before mutation and always restore in finally.

6. Public integration and docs/examples are missing.

This adds a public MultiLoRA class and lifecycle helpers, but config-driven PEFT cannot instantiate it because create_peft() still only maps lora, vlm_lora, canonical_lora, and dora. The PR also does not update docs/training/peft.md or add an example showing how users should initialize slots, route tokens, load/clear adapters, and export one slot.

Impact: users cannot discover or use this feature through normal config/recipe paths, and the required lifecycle is easy to misuse.

Suggested fix: add type: multi_lora support if this is meant to be user-facing, update PEFT docs, and add a minimal example covering slot init, set_tokens_per_adapter_slot, checkpoint/load/export lifecycle, and current limitations.

[mathewjhan]

@yaoyu-33 thanks for the detailed review, super helpful! I will address these and report back

[yaoyu-33]

@mathewjhan I asked agent to do a pass first, feel free to ignore any items that not make sense. I am going to do another round later.