bridge.peft.base#

Module Contents#

Classes#

PEFT

Abstract base class for Parameter-Efficient Fine-Tuning (PEFT) methods.

Data#

logger

ModelType

API#

bridge.peft.base.logger: logging.Logger#

‘getLogger(…)’

bridge.peft.base.ModelType#

‘TypeVar(…)’

class bridge.peft.base.PEFT#

Bases: abc.ABC

Abstract base class for Parameter-Efficient Fine-Tuning (PEFT) methods.

This class defines the interface for PEFT methods, which are used to fine-tune large language models efficiently by modifying only a small subset of the model’s parameters.

.. rubric:: Example

class MyPEFT(PEFT): def transform(self, module, name=None, prefix=None): # Implement the transform logic pass

from megatron.bridge.models import get_base_model

peft = MyPEFT() base_model = get_base_model(model_config) # Returns list[MegatronModule] adapted_model = peft(base_model, training=True)

params_to_save: set[str]#

‘field(…)’

abstractmethod transform(
module: torch.nn.Module,
name: Optional[str] = None,
prefix: Optional[str] = None,
) → torch.nn.Module#

Transform a single module according to the PEFT method.

This method is called for each module in the model during the PEFT application process. It should be implemented by subclasses to define how individual modules are transformed for the specific PEFT technique.

Parameters:

  • module (nn.Module) – The individual module to be transformed.

  • name (Optional[str]) – The name of the module within the model structure. Defaults to None.

  • prefix (Optional[str]) – A prefix to be added to the module name, typically used for nested modules. Defaults to None.

Returns:

The transformed module. This can be the original module with modifications, a new module replacing the original, or the original module if no transformation is needed for this specific module.

Return type:

nn.Module

.. note::

This method is automatically called for each module in the model when the PEFT instance is applied to the model using the call method.

__call__(
model: bridge.peft.base.ModelType,
training: bool = True,
) → bridge.peft.base.ModelType#

Apply the PEFT method to the entire model.

This method freezes the model parameters and walks through the model structure, applying the transform method to each module.

Parameters:

  • model – The model to be fine-tuned. Can be a single model or a list of model chunks (for pipeline parallelism).

  • training (bool) – Whether the model will be used for training. If False, additional freezing may be applied. Defaults to True.

Returns:

The same type as the input model, transformed with PEFT applied.

freeze_model(
model: bridge.peft.base.ModelType,
training: bool = True,
) → None#

Apply a default freeze method to the model.

This method freezes all the model parameters. This method can be overridden by subclasses to implement custom freeze strategies (e.g. freeze only parts of the model)

Parameters:

  • model – The model to be fine-tuned.

  • training (bool) – Whether the model is being used for training. Affects training mode handling.

set_params_to_save(model: bridge.peft.base.ModelType) → None#

Set parameters to be saved for PEFT checkpointing.

This method identifies which parameters should be saved during checkpointing to reduce storage requirements (only adapter parameters, not the full model).

Parameters:

model – The model after PEFT has been applied.

adapter_key_filter(key) → bool#

Filter function for adapter parameters during checkpointing.

This method determines if a parameter should be included in checkpoints. Used to save only adapter weights, not the full model.

Parameters:

key (str or tuple) – Parameter name/key to check. Can be a string for regular checkpointing or a tuple for distributed checkpointing.

Returns:

True if the parameter should be saved.

Return type:

bool