Implementation:Hpcaitech ColossalAI LoRA Module

Knowledge Sources	Hpcaitech_ColossalAI
Domains	Parameter_Efficient_Finetuning, LoRA, Model_Adaptation
Last Updated	2026-02-09 00:00 GMT

Overview

lora.py implements Low-Rank Adaptation (LoRA) for efficient fine-tuning, providing LoRA-wrapped linear and embedding layers with support for both Kaiming uniform and PiSSA (SVD-based) initialization.

Description

This module provides a complete LoRA implementation with several key components. LoraConfig is a dataclass for configuring rank, alpha, dropout, target modules, and initialization method. LoraBase is the abstract base class inheriting from both lora.LoRALayer and nn.Module, implementing weight merging/unmerging logic for eval/train mode transitions. LoraLinear wraps nn.Linear layers with low-rank adapters (lora_A, lora_B) using out-of-place operations compatible with ColossalAI Gemini. LoraEmbedding similarly wraps nn.Embedding layers. Both support PiSSA initialization (Principal Singular values and Singular vectors Adaptation) via SVD decomposition, where the frozen weight stores the less significant components and the trainable adapters are initialized from the principal components. The convert_to_lora_module function recursively converts all matching linear and embedding layers in a model to LoRA variants, freezing original parameters and logging trainable parameter statistics.

Usage

Use this module to convert any PyTorch model to a parameter-efficient LoRA variant for fine-tuning. Apply convert_to_lora_module to a pre-trained model before training to freeze most parameters and only train the low-rank adapter matrices.

Code Reference

Source Location

• Repository: Hpcaitech_ColossalAI
• File: applications/ColossalChat/coati/models/lora.py
• Lines: 1-367

Signature

@dataclasses.dataclass
class LoraConfig:
    r: int = 0
    lora_alpha: int = 32
    linear_lora_dropout: float = 0.1
    embedding_lora_dropout: float = 0.0
    lora_train_bias: str = "none"
    lora_initialization_method: str = "kaiming_uniform"
    target_modules: List = None

class LoraLinear(LoraBase):
    def __init__(
        self,
        weight: nn.Parameter,
        bias: Union[nn.Parameter, bool],
        r: int = 0,
        lora_alpha: int = 32,
        lora_dropout: float = 0.0,
        lora_initialization_method: str = "kaiming_uniform",
    )

class LoraEmbedding(LoraBase):
    def __init__(
        self,
        weight: nn.Parameter,
        r: int = 0,
        lora_alpha: int = 32,
        lora_dropout: float = 0.1,
        num_embeddings: int = None,
        embedding_dim: int = None,
        padding_idx: Optional[int] = None,
        max_norm: Optional[float] = None,
        norm_type: float = 2.0,
        scale_grad_by_freq: bool = False,
        sparse: bool = False,
        lora_initialization_method: str = "kaiming_uniform",
    )

def convert_to_lora_module(module: nn.Module, lora_config: LoraConfig) -> nn.Module

Import

from coati.models.lora import LoraConfig, LoraLinear, LoraEmbedding, convert_to_lora_module

I/O Contract

Inputs (convert_to_lora_module)

Name	Type	Required	Description
module	nn.Module	Yes	The PyTorch model to convert to LoRA
lora_config	LoraConfig	Yes	Configuration specifying rank, alpha, dropout, target modules, and initialization method

Outputs (convert_to_lora_module)

Name	Type	Description
module	nn.Module	The converted model with LoRA layers (original parameters frozen, adapters trainable)

LoraConfig Parameters

Name	Type	Default	Description
r	int	0	LoRA rank (0 disables LoRA)
lora_alpha	int	32	LoRA scaling factor
linear_lora_dropout	float	0.1	Dropout for linear LoRA layers
embedding_lora_dropout	float	0.0	Dropout for embedding LoRA layers
lora_train_bias	str	"none"	Bias training mode: "none", "all", or "lora"
lora_initialization_method	str	"kaiming_uniform"	Initialization: "kaiming_uniform" or "PiSSA"
target_modules	List	None	List of module name patterns to convert (None converts all)

Usage Examples

from coati.models.lora import LoraConfig, convert_to_lora_module
from transformers import AutoModelForCausalLM

# Load a pre-trained model
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# Configure LoRA
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    linear_lora_dropout=0.1,
    lora_train_bias="none",
    lora_initialization_method="kaiming_uniform",
    target_modules=["q_proj", "v_proj"],
)

# Convert model to LoRA
model = convert_to_lora_module(model, lora_config)
# Logs: Trainable parameter size: X.XXM, Percentage: X.XX%

# Using PiSSA initialization
lora_config_pissa = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_initialization_method="PiSSA",
)
model_pissa = convert_to_lora_module(model, lora_config_pissa)

# Load config from JSON file
lora_config = LoraConfig.from_file("lora_config.json")

Related Pages

• Environment:Hpcaitech_ColossalAI_GRPO_Distributed_Environment