Method default_config

axlearn/common/bert.py:325–348

(cls)

    @classmethod
    def default_config(cls):
        cfg: BertModel.Config = super().default_config()
        # param_init matches original tf implementation. The Hugging Face implementation uses plain
        # GaussianInitializer which empirically does not work as well (gradients quickly vanish).
        # https://github.com/google-research/bert/blob/eedf5716ce1268e56f0a50264a88cafad334ac61/modeling.py#L377
        # https://github.com/huggingface/transformers/blob/aa6cfe9c4b073b2c058a78fc2d26fe3fbe0ad70b/src/transformers/models/bert/modeling_bert.py#L726
        cfg.param_init = DefaultInitializer.default_config().set(
            init_by_param_name={
                PARAM_REGEXP_WEIGHT: WeightInitializer.default_config().set(
                    fan=None,
                    distribution="truncated_normal",
                    scale=0.02,
                )
            }
        )
        # By default, assume `head` employs tied weights.
        cfg.head = BertLMHead.default_config().set(
            inner_head=RedirectToSharedModule.default_config().set(
                shared_module="shared_token_emb",
                # Map the method name here so the child can just call forward directly.
                method_map=dict(forward="attend"),
            )
        )
        return cfg


def bert_embedding_config(