| 47 | |
| 48 | |
| 49 | class Lora(Module): |
| 50 | |
| 51 | def __init__(self, |
| 52 | in_hidden_size: int = 0, |
| 53 | out_hidden_sizes: List[int] = [0], |
| 54 | max_low_rank: int = 0) -> None: |
| 55 | super().__init__() |
| 56 | |
| 57 | self.in_hidden_size = in_hidden_size |
| 58 | self.out_hidden_sizes = out_hidden_sizes |
| 59 | self.max_low_rank = max_low_rank |
| 60 | |
| 61 | def forward(self, |
| 62 | x, |
| 63 | lora_runtime_params: LoraRuntimeParams = None, |
| 64 | is_cross_attention: bool = False): |
| 65 | if default_net().plugin_config.lora_plugin: |
| 66 | result = lora_plugin( |
| 67 | x, |
| 68 | in_hidden_size=self.in_hidden_size, |
| 69 | out_hidden_sizes=self.out_hidden_sizes, |
| 70 | host_request_types=lora_runtime_params.host_request_types, |
| 71 | transb=True, |
| 72 | # For cross attention, host_encoder_input_lengths should be used instead of host_context_lengths |
| 73 | host_context_lengths=lora_runtime_params.host_context_lengths |
| 74 | if not is_cross_attention else |
| 75 | lora_runtime_params.host_encoder_input_lengths, |
| 76 | max_low_rank=self.max_low_rank, |
| 77 | lora_ranks=lora_runtime_params.lora_ranks, |
| 78 | lora_weights_pointers=lora_runtime_params.lora_weights_pointers, |
| 79 | weight_index=lora_runtime_params.weight_index, |
| 80 | ) |
| 81 | if lora_runtime_params.partial_lora_mask is not None: |
| 82 | zero_tensor = constant(np.array([0.0], dtype=np.float16)) |
| 83 | if isinstance(result, List): |
| 84 | result = [ |
| 85 | where(lora_runtime_params.partial_lora_mask, r, |
| 86 | zero_tensor) for r in result |
| 87 | ] |
| 88 | elif isinstance(result, Tensor): |
| 89 | result = where(lora_runtime_params.partial_lora_mask, |
| 90 | result, zero_tensor) |
| 91 | else: |
| 92 | assert False |
| 93 | else: |
| 94 | assert False, "Not support lora without plugin" |
| 95 | |
| 96 | return result |
| 97 | |
| 98 | |
| 99 | class Dora(Module): |
no outgoing calls