Method __init__

axlearn/common/lora.py:287–311

(self, cfg: Config, *, parent: Module)

            )

    def __init__(self, cfg: Config, *, parent: Module):
        super().__init__(cfg, parent=parent)
        cfg = self.config
        self._add_child(
            "lora_down",
            cfg.lora_down.set(
                model_dim=cfg.rank,
                num_heads=cfg.num_heads,
                per_head_dim=self._per_head_dim,
                bias=False,
            ),
        )
        self._add_child(
            "lora_up",
            cfg.lora_up.set(
                input_dim=cfg.rank,
                output_dim=cfg.output_dim,
                bias=False,
                param_init=DefaultInitializer.default_config().set(
                    init_by_param_name={
                        PARAM_REGEXP_WEIGHT: ConstantInitializer.default_config().set(value=0.0)
                    }
                ),
            ),
        )


class LoraFusedQKVAdapter(_BaseLoraAdapter):