Skip to content

vllm_gaudi.lora.layers.hpu_row_parallel_linear

HPU-specific LoRA layer for RowParallelLinear.

This module registers HPU-aware LoRA wrappers that can handle HPURowParallelLinear instances created via OOT registration.

HPURowParallelLinearWithLoRA

Bases: RowParallelLinearWithLoRA

LoRA wrapper that can handle HPURowParallelLinear instances.

Source code in vllm_gaudi/lora/layers/hpu_row_parallel_linear.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class HPURowParallelLinearWithLoRA(RowParallelLinearWithLoRA):
    """LoRA wrapper that can handle HPURowParallelLinear instances."""

    @classmethod
    @_not_fully_sharded_can_replace
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: list,
        model_config: PretrainedConfig | None = None,
    ) -> bool:
        # Accept both RowParallelLinear and HPURowParallelLinear
        return isinstance(source_layer, RowParallelLinear)

can_replace_layer classmethod 

can_replace_layer(
    source_layer: Module,
    lora_config: LoRAConfig,
    packed_modules_list: list,
    model_config: PretrainedConfig | None = None,
) -> bool
Source code in vllm_gaudi/lora/layers/hpu_row_parallel_linear.py 
28
29
30
31
32
33
34
35
36
37
38
@classmethod
@_not_fully_sharded_can_replace
def can_replace_layer(
    cls,
    source_layer: nn.Module,
    lora_config: LoRAConfig,
    packed_modules_list: list,
    model_config: PretrainedConfig | None = None,
) -> bool:
    # Accept both RowParallelLinear and HPURowParallelLinear
    return isinstance(source_layer, RowParallelLinear)

HPURowParallelLinearWithShardedLoRA

Bases: RowParallelLinearWithShardedLoRA

Sharded LoRA wrapper that can handle HPURowParallelLinear instances.

Source code in vllm_gaudi/lora/layers/hpu_row_parallel_linear.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class HPURowParallelLinearWithShardedLoRA(RowParallelLinearWithShardedLoRA):
    """Sharded LoRA wrapper that can handle HPURowParallelLinear instances."""

    @classmethod
    @_fully_sharded_can_replace
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: list,
        model_config: PretrainedConfig | None = None,
    ) -> bool:
        # Use isinstance to accept HPURowParallelLinear (subclass of RowParallelLinear)
        return isinstance(source_layer, RowParallelLinear)

can_replace_layer classmethod 

can_replace_layer(
    source_layer: Module,
    lora_config: LoRAConfig,
    packed_modules_list: list,
    model_config: PretrainedConfig | None = None,
) -> bool
Source code in vllm_gaudi/lora/layers/hpu_row_parallel_linear.py 
44
45
46
47
48
49
50
51
52
53
54
@classmethod
@_fully_sharded_can_replace
def can_replace_layer(
    cls,
    source_layer: nn.Module,
    lora_config: LoRAConfig,
    packed_modules_list: list,
    model_config: PretrainedConfig | None = None,
) -> bool:
    # Use isinstance to accept HPURowParallelLinear (subclass of RowParallelLinear)
    return isinstance(source_layer, RowParallelLinear)

register_hpu_lora_layers 

register_hpu_lora_layers()

Register HPU LoRA layers in the global lora_utils._all_lora_classes set.

This must be called before LoRA model creation to ensure HPU layers are checked before upstream layers.

Source code in vllm_gaudi/lora/layers/hpu_row_parallel_linear.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
def register_hpu_lora_layers():
    """Register HPU LoRA layers in the global lora_utils._all_lora_classes set.

    This must be called before LoRA model creation to ensure HPU layers
    are checked before upstream layers.
    """
    # Remove the upstream classes first (they use strict type check)
    lora_utils._all_lora_classes.discard(RowParallelLinearWithLoRA)
    lora_utils._all_lora_classes.discard(RowParallelLinearWithShardedLoRA)

    # Add HPU-aware classes
    lora_utils._all_lora_classes.add(HPURowParallelLinearWithLoRA)
    lora_utils._all_lora_classes.add(HPURowParallelLinearWithShardedLoRA)