Skip to content

peft.md

Latest commit

 

History

History
460 lines (362 loc) · 24.1 KB

peft.md

File metadata and controls

460 lines (362 loc) · 24.1 KB
Raw

Parameter-Efficient Fine-Tuning (PEFT)

This guide explains how to configure and use PEFT in Megatron Bridge—covering LoRA and DoRA, required checkpoints, example configurations, and the internal design and training workflow—so you can integrate, scale, and checkpoint adapters efficiently.

Model Customization

Customizing models enables you to adapt a general pre-trained model to a specific use case or domain. This process produces a fine-tuned model that retains the broad knowledge from pretraining while delivering more accurate outputs for targeted downstream tasks.

Model customization is typically achieved through supervised fine-tuning, which falls into two main approaches: Full-Parameter Fine-Tuning, known as Supervised Fine-Tuning (SFT), and Parameter-Efficient Fine-Tuning (PEFT).

In SFT, all model parameters are updated to align the model’s outputs with the task-specific requirements. This approach often yields the highest performance but can be computationally intensive.

PEFT, by contrast, updates only a small subset of parameters that are inserted into the base model at strategic locations. The base model weights remain frozen, and only the adapter modules are trained. This significantly reduces the number of trainable parameters—often to less than 1%—while still achieving near-SFT levels of accuracy.

As language models continue to grow in size, PEFT is gaining popularity for its efficiency and minimal hardware demands, making it a practical choice for many real-world applications.

PEFT Configuration

PEFT is configured as an optional attribute in ConfigContainer:

from megatron.bridge.peft.lora import LoRA
from megatron.bridge.training.config import CheckpointConfig, ConfigContainer

config = ConfigContainer(
    # ... other required configurations
    peft=LoRA(
        target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"],
        dim=16,
        alpha=32,
        dropout=0.1,
    ),
    checkpoint=CheckpointConfig(
        pretrained_checkpoint="/checkpoints/base_model",  # Required for PEFT
        save="/path/to/peft/checkpoints",
    ),
)
**Requirement:** PEFT requires `checkpoint.pretrained_checkpoint` to load the frozen base model weights before adapters are inserted. This path may be a native Megatron checkpoint base directory, a specific native Megatron `iter_N` directory, or a local Hugging Face full-model directory containing `config.json` and model weight files. A remote Hugging Face model ID is not a checkpoint path; download it locally or convert it to Megatron format first.

When resuming adapter training, keep `checkpoint.pretrained_checkpoint` pointed at the same base model and set `checkpoint.load` to the PEFT adapter checkpoint directory.

MoE router state and checkpoint compatibility

When expert-bias routing is enabled, gradient finalization updates the router's expert_bias buffer even when the router weights are frozen. PEFT checkpoints therefore retain these buffers alongside adapter parameters. Adapter reload and merge into a full model preserve this state.

Older adapter checkpoints may omit these buffers. The adapter inference/merge loader rejects such checkpoints for a bias-enabled model, including when its strict=False option is used: that option does not permit missing distributed checkpoint tensors. The separate training-resume loader follows checkpoint.dist_ckpt_strictness; a permissive policy may retain the base-model bias, which is not exact recovery of the trained router state. Missing trained biases cannot be reconstructed from adapter weights alone.

Standalone Hugging Face adapter export does not serialize router-bias buffers. For models whose biases changed during PEFT, merge the native adapter checkpoint into the base model and export the resulting full Hugging Face model instead.

Supported PEFT Methods

LoRA: Low-Rank Adaptation of Large Language Models

LoRA makes fine-tuning efficient by representing weight updates with two low-rank decomposition matrices. The original model weights remain frozen, while the low-rank decomposition matrices are updated to adapt to the new data, keeping the number of trainable parameters low. In contrast with adapters, the original model weights and adapted weights can be combined during inference, avoiding any architectural change or additional latency in the model at inference time.

In Megatron Bridge, you can configure both the adapter bottleneck dimension and the target modules where LoRA is applied. LoRA supports any linear layer, which in transformer models typically includes:

  1. Query, key, and value (QKV) attention projections
  2. The attention output projection
  3. One or both MLP layers

Megatron Bridge fuses the QKV projections into a single linear layer. As a result, LoRA learns a unified low-rank adaptation for the combined QKV representation.

from megatron.bridge.peft.lora import LoRA

lora_config = LoRA(
    target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"],
    dim=16,                    # Rank of adaptation
    alpha=32,                  # Scaling parameter  
    dropout=0.1,               # Dropout rate
)

Key Parameters

The following table lists key hyperparameters for configuring LoRA, which control its module targeting, adaptation rank, scaling behavior, and regularization strategy.

Parameter Type Default Description
target_modules List[str] All linear layers Modules to apply LoRA to
dim int 32 Rank of the low-rank adaptation
alpha float 32 Scaling parameter for LoRA
dropout float 0.0 Dropout rate for LoRA layers
normalize_moe_lora bool False Reduce expert-layer rank to dim // moe_router_topk while keeping dense layers at the full rank
share_expert_adapters bool True Share one adapter across all local experts on an EP rank instead of creating one adapter per local expert

Target Modules

The following table lists specific submodules within transformer architectures that are commonly targeted for LoRA, enabling efficient fine-tuning of attention and feedforward components:

Module Description
linear_qkv Query, key, value projections in attention
linear_proj Attention output projection
linear_fc1 First MLP layer
linear_fc2 Second MLP layer

Wildcard Target Modules

For more granular targeting, individual layers can be targeted for the adapters.

# Target specific layers only
lora_config = LoRA(
    target_modules=[
        "*.layers.0.*.linear_qkv",   # First layer only
        "*.layers.1.*.linear_qkv",   # Second layer only
    ]
)

MoE Expert Adapter Modes

LoRA and CanonicalLoRA expose MoE-specific controls for grouped expert MLP layers such as linear_fc1, linear_fc2, linear_fc1_up, and linear_fc1_gate:

Parameter Type Default Description
share_expert_adapters bool True Preserve the original behavior and share one adapter across all local experts on each EP rank
normalize_moe_lora bool False Use dim // moe_router_topk for expert layers only so MoE adapter capacity stays comparable to a dense model
experts_shared_outer_loras bool False LoRA only. Shared-outer layout for grouped experts: the gate/up LoRA-A and the down LoRA-B are shared across experts while the other factor stays per expert (SGLang's experts_shared_outer_loras serving contract)
  • share_expert_adapters=True keeps the original shared-adapter behavior for grouped expert linears.
  • share_expert_adapters=False opts into per-local-expert adapters.
  • normalize_moe_lora=True only changes expert layers; non-expert layers still use the full dim.
  • dim must be divisible by moe_router_topk when normalize_moe_lora=True.
  • When expert tensor parallelism shards a non-input-parallel expert layer, Megatron Bridge may round the effective expert rank up to the expert-TP granularity so the adapter can be partitioned cleanly.
from megatron.bridge.peft.lora import LoRA

lora_config = LoRA(
    target_modules=["linear_fc1", "linear_fc2"],
    dim=8,
    normalize_moe_lora=True,
    share_expert_adapters=False,
)
Expert-parallel synchronization of shared adapters

Adapter weights that serve every expert (the default shared adapter, and the shared factor of each experts_shared_outer_loras=True adapter) are replicated across expert-parallel ranks, but Megatron-Core's expert DDP only reduces gradients over expert-data-parallel replicas. Megatron Bridge therefore sums those gradients across the expert-parallel group once per step, coalesced into a single collective right after the data-parallel gradient sync, the same point where Megatron-Core sums sequence-parallel layernorm gradients across tensor parallelism. The training loop installs this automatically for PEFT runs through finalize_model_grads_with_expert_adapter_sync (megatron.bridge.peft.utils), and it stays correct with gradient_accumulation_fusion enabled because it reads the fused main_grad buffers.

External training loops that call Megatron-Core's finalize_model_grads themselves must install the Bridge wrapper (or call allreduce_expert_parallel_replicated_grads right after their own finalize step):

from functools import partial

from megatron.bridge.peft.utils import finalize_model_grads_with_expert_adapter_sync

model_config.finalize_model_grads_func = partial(
    finalize_model_grads_with_expert_adapter_sync, pg_collection=pg_collection
)

A per-layer fallback hook keeps eagerly computed weight gradients in sync until the finalize path takes over (Bridge's setup removes the hooks before the first step; otherwise the wrapper's first call does). The fallback cannot cover fused weight-gradient accumulation and logs a warning in that case.

Exporting shared-outer MoE adapters

AutoBridge.export_adapter_weights emits shared-outer adapters in the legacy (SGLang) layout unless a serving layout is requested: the shared factor once as a [1, ...] tensor under the expert-agnostic name, plus per-expert factors under experts.N.* names. Three opt-in flags select a serving layout instead, listed from highest to lowest priority:

Option Default Export layout when enabled
stack_3d_moe False vLLM 3D-MoE layout: stacked tensors under experts.base_layer (gate/up) and experts (down)
moe_shared_loras False vLLM 2D shared-MoE layout under experts.w1/w2/w3, with the shared factors kept at expert-dim 1
expand_shared_outer False vLLM 2D pack_moe layout: the shared factor replicated under per-expert experts.N.* names

The flags can be combined; when several are set, the highest-priority one decides the layout. With all three unset, existing callers keep the legacy layout. moe_shared_loras=True is the 2D serving format used when vLLM's enable_moe_shared_loras setting is enabled; it does not imply a 3D model. Callers can select stack_3d_moe from the model's MoE tensor layout and moe_shared_loras from the serving backend's adapter mode.

LoRA+

LoRA+ trains the B matrix (linear_out) at a higher learning rate than A (linear_in) by a fixed ratio lr_B = lora_plus_ratio * lr_A (paper default 16). The ratio is set on the optimizer, not the LoRA adapter: when it differs from 1.0, get_lora_plus_config_overrides (megatron.bridge.peft.lora) builds Megatron-Core per-group overrides giving the *.linear_out.weight group max_lr/min_lr scaled by the ratio. Pass 1.0 (or skip the call) to keep a single learning rate.

from megatron.bridge.peft.lora import get_lora_plus_config_overrides
from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer

opt_config = OptimizerConfig(optimizer="adam", lr=1e-4, bf16=True)

# LoRA+ LR split: lr_B = lora_plus_ratio * lr_A; pass 1.0 (or skip) to disable.
lora_plus_ratio = 16
config_overrides = get_lora_plus_config_overrides(opt_config, lora_plus_ratio)
optimizer = get_megatron_optimizer(
    config=opt_config,
    model_chunks=model_chunks,
    config_overrides=config_overrides,
)

Canonical LoRA: Performant vs Canonical Variants

There are two variants of LoRA implemented in Megatron Bridge: "performant LoRA" (LoRA) and "canonical LoRA" (CanonicalLoRA).

The distinction comes from the fact that Megatron Core optimizes the implementation of the following two linear modules by fusing multiple linear layers into one layer. When these layers are adapted with LoRA, the performant version also uses only one adapter for the linear module. The two linear modules are:

  1. linear_qkv: The projection matrix in self attention that transforms hidden state to query, key and value. Megatron Core fuses these three projection matrices into a single matrix to efficiently parallelize the matrix multiplication. Hence, performant LoRA applies a single adapter to the qkv projection matrix, whereas canonical LoRA applies three adapters.
  2. linear_fc1: The first linear layer in the MLP module before the intermediate activation. For gated linear activations, Megatron Core fuses the up and gate projection matrices into a single matrix for efficient parallelization. Hence, performant LoRA applies a single adapter to the up and gate projection matrices, whereas canonical LoRA applies two adapters.

The following two figures illustrate the difference between canonical and performant LoRA, using the linear_qkv layer as an example. Canonical LoRA runs three adapters sequentially, while performant LoRA runs one adapter.

:width: 640
:align: center

:width: 400
:align: center

Canonical LoRA conforms more closely to reference implementations, though it is slower in comparison since it performs several matrix multiplications sequentially, as described above. Performant LoRA has fewer parameters than canonical LoRA and can often achieve the same level of accuracy as canonical LoRA.

Though not immediately apparent, performant LoRA is mathematically equivalent to canonical LoRA when the $A_q$, $A_k$, $A_v$ matrices are tied (i.e. forced to share the same weight during training) in linear_qkv, and similarly when the $A_{up}$, $A_{gate}$ matrices are tied in linear_fc1.

:class: dropdown

Let $[x \quad y]$ denote matrix concatenation. (In Megatron Bridge, this concatenation is done in an interleaved fashion, but this does not affect the proof below.)

Let $A_q = A_k = A_v = A_{qkv}$ (weight tying)

Then

$$
\begin{aligned}
[query \quad key \quad value]
&= [W_q x + B_q A_q x \quad W_k x + B_k A_k x \quad W_v x + B_v A_v x] \quad\quad \text{(canonical formulation)} \\
&= [W_q x + B_q (A_{qkv} x) \quad W_k x + B_k (A_{qkv} x) \quad W_v x + B_v (A_{qkv} x)] \\
&= [W_q \quad W_k \quad W_v] x + [B_q \quad B_k \quad B_v]A_{qkv} x \\
&= W_{qkv} x + B_{qkv} A_{qkv} x  \quad\quad \text{(performant formulation)}
\end{aligned}
$$

Note: dimensions of weight matrices are as follows:

$$
\begin{align}
W_q:     &\ n_q d \times h          \qquad & A_q:     &\ r \times h \qquad  & B_q:     &\ n_q d \times r \\
W_k:     &\ n_{kv} d \times h       \qquad & A_k:     &\ r \times h \qquad  & B_k:     &\ n_{kv} d \times r \\
W_v:     &\ n_{kv} d \times h       \qquad & A_v:     &\ r \times h \qquad  & B_v:     &\ n_{kv} d \times r \\
W_{qkv}: &\ (n_q+2n_{kv})d \times h \qquad & A_{qkv}: &\ r \times h \qquad  & B_{qkv}: &\ (n_q+2n_{kv})d \times r
\end{align}
$$

Where:
- $n_q$: Number of attention heads (`num_attention_heads`).
- $n_{kv}$: Number of key value heads (`num_query_groups`). Note that if grouped query attention (GQA) is not used, $n_{kv} = n_q$.
- $h$: Transformer hidden size (`hidden_size`).
- $d$: Transformer head dimension (`kv_channels`).
- $r$: LoRA rank.

Using Canonical LoRA

from megatron.bridge.peft.canonical_lora import CanonicalLoRA

canonical_lora_config = CanonicalLoRA(
    target_modules=[
        "linear_q", "linear_k", "linear_v",      # Individual Q, K, V projections
        "linear_proj",                           # Attention output projection
        "linear_fc1_up", "linear_fc1_gate",     # Individual up and gate projections
        "linear_fc2"                             # Second MLP layer
    ],
    dim=16,                    # Rank of adaptation
    alpha=32,                  # Scaling parameter
    dropout=0.1,               # Dropout rate
)

CanonicalLoRA exposes the same normalize_moe_lora and share_expert_adapters controls as LoRA for MoE expert layers.

Key Parameters

Parameter Type Default Description
target_modules List[str] All canonical linear layers Modules to apply canonical LoRA to
dim int 32 Rank of the low-rank adaptation
alpha float 32 Scaling parameter for LoRA
dropout float 0.0 Dropout rate for LoRA layers
dropout_position Literal["pre", "post"] "pre" Position for applying dropout
lora_A_init_method str "xavier" Initialization method for LoRA A matrix
lora_B_init_method str "zero" Initialization method for LoRA B matrix
normalize_moe_lora bool False Reduce expert-layer rank to dim // moe_router_topk while keeping dense layers at the full rank
share_expert_adapters bool True Share one adapter across all local experts on an EP rank instead of creating one adapter per local expert

Target Modules for Canonical LoRA

The following table lists specific submodules within transformer architectures that are targeted for canonical LoRA:

Module Description
linear_q Query projection in attention
linear_k Key projection in attention
linear_v Value projection in attention
linear_proj Attention output projection
linear_fc1_up Up projection in MLP
linear_fc1_gate Gate projection in MLP
linear_fc2 Second MLP layer 
Canonical LoRA does not support `linear_qkv` or `linear_fc1` targets. Use the individual component targets (`linear_q`, `linear_k`, `linear_v` for QKV and `linear_fc1_up`, `linear_fc1_gate` for FC1) instead.

DoRA: Weight-Decomposed Low-Rank Adaptation

DoRA decomposes the pre-trained weight into magnitude and direction. It learns a separate magnitude parameter while employing LoRA for directional updates, efficiently minimizing the number of trainable parameters. DoRA enhances both the learning capacity and training stability of LoRA, while avoiding any additional inference overhead. DoRA has been shown to consistently outperform LoRA on various downstream tasks.

In Megatron Bridge, DoRA leverages the same adapter structure as LoRA. Megatron Bridge adds support for Tensor Parallelism and Pipeline Parallelism for DoRA, enabling DoRA to be scaled to larger model variants.

from megatron.bridge.peft.dora import DoRA

dora_config = DoRA(
    target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"],
    dim=16,                    # Rank of adaptation
    alpha=32,                  # Scaling parameter
    dropout=0.1,               # Dropout rate
)

Key Parameters

The following parameters define how LoRA is applied to your model. They control which modules are targeted, the adaptation rank, scaling behavior, and dropout configuration:

Parameter Type Default Description
target_modules List[str] All linear layers Modules to apply DoRA to
dim int 32 Rank of the low-rank adaptation
alpha float 64 Scaling parameter for DoRA
dropout float 0.0 Dropout rate for DoRA layers

Full Configuration Example

from megatron.bridge.training.config import (
    CheckpointConfig,
    ConfigContainer,
    SchedulerConfig,
    TrainingConfig,
)
from megatron.bridge.data.builders import (
    GPTSFTDatasetConfig,
    HFDatasetSourceConfig,
    PromptCompletionSFTPreprocessingConfig,
)
from megatron.bridge.data.packing import PackedSequenceSpecs
from megatron.bridge.peft.lora import LoRA
from megatron.core.optimizer import OptimizerConfig

# Configure PEFT fine-tuning
config = ConfigContainer(
    model=model_provider,
    train=TrainingConfig(
        train_iters=1000,
        global_batch_size=64,
        micro_batch_size=1,  # Required for packed sequences if used
        eval_interval=100,
    ),
    optimizer=OptimizerConfig(
        optimizer="adam",
        lr=1e-4,  # Lower learning rate for fine-tuning
        weight_decay=0.01,
        bf16=True,
        use_distributed_optimizer=True,
    ),
    scheduler=SchedulerConfig(
        lr_decay_style="cosine",
        lr_warmup_iters=100,
        lr_decay_iters=1000,
    ),
    dataset=GPTSFTDatasetConfig(
        seq_length=512,
        hf_dataset=HFDatasetSourceConfig(dataset_name="squad"),
        preprocessing=PromptCompletionSFTPreprocessingConfig(separator=" "),
        hf_validation_proportion=0.1,
        seed=5678,
        do_validation=True,
        do_test=False,
        dataset_kwargs={"pad_to_max_length": True},
        enable_offline_packing=True,
        offline_packing_specs=PackedSequenceSpecs(packed_sequence_size=512),
    ),
    checkpoint=CheckpointConfig(
        pretrained_checkpoint="/checkpoints/base_model",  # Required for PEFT
        save="/path/to/peft/checkpoints",
        save_interval=200,
    ),
    peft=LoRA(
        target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"],
        dim=16,
        alpha=32,
        dropout=0.1,
    ),
    # ... other configurations
)

PEFT Design in Megatron Bridge

This section describes the internal design and architecture for how PEFT is integrated into Megatron Bridge.

Architecture Overview

The PEFT framework introduces a modular design for integrating adapters into large-scale models. Its architecture consists of the following components:

  1. Base PEFT Class: All PEFT methods inherit from the abstract {py:class}bridge.peft.base.PEFT base class, which defines the core interface for module transformation.
  2. Module Transformation: PEFT traverses the model structure to identify and transform target modules individually.
  3. Adapter Integration: Adapters are injected into selected modules using a pre-wrap hook during model initialization.
  4. Checkpoint Integration: Adapter parameters and updated MoE router expert-bias buffers are saved and loaded during checkpointing; base model weights remain frozen.

PEFT Workflow in Training

The training workflow for PEFT follows a structured sequence that ensures efficient fine-tuning with minimal overhead:

  1. Model Loading: The base model is initialized from a specified pretrained checkpoint.
  2. PEFT Application: Adapter transformations are applied after Megatron Core model initialization, but before distributed wrapping.
  3. Parameter Freezing: Base model parameters are frozen to reduce training complexity; only adapter parameters are updated.
  4. Adapter Weight Loading: When resuming training, adapter weights and saved router expert-bias buffers are restored from the checkpoint.
  5. Checkpoint Saving: Adapter states and router expert-bias buffers are saved, resulting in significantly smaller checkpoint files than full-model checkpoints.

Key Benefits

PEFT offers several advantages for scalable and efficient model fine-tuning:

  • Reduced Checkpoint Size: Adapter-only checkpoints are dramatically smaller than full model checkpoints.
  • Memory Efficiency: Since gradients are computed only for adapter parameters, memory usage is significantly reduced.
  • Resume Support: Training can be resumed seamlessly using adapter-only checkpoints, without reloading full model weights.