Skip to content

Consolidate PEFT support into Megatron-LM #7267

Description

@maanug-nv

Goal

Move the PEFT support required by Megatron-Bridge into Megatron-LM so Bridge PEFT recipes can use megatron.training as the canonical training loop.

Scope

  • Move performant LoRA, canonical LoRA, DoRA linear adapters, and the LoRALinear wrapper into Megatron Core.
  • Apply adapters and freeze the base model before model wrapping in the Megatron-LM training loop.
  • Add adapter-only checkpoint save and load support to megatron.training.checkpointing.
  • Migrate Megatron-Bridge PEFT recipes to the Megatron-LM training loop.

Completion criteria

  • Existing Megatron-Bridge PEFT APIs continue to work through the migration.
  • Existing PEFT recipes retain their current behavior when run through the Megatron-LM training loop.
  • Representative performant LoRA, canonical LoRA, and DoRA recipes preserve convergence and performance within agreed tolerances.
  • Adapter-only checkpointing, adapter merging, and adapter export continue to work.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions