fix: try to get from_config to work

modeling_bert.py

@@ -336,7 +336,7 @@ class BertPreTrainedModel(nn.Module):
         return model
 
     @classmethod
-    def from_config(cls, config, *inputs, **kwargs):
+    def _from_config(cls, config, *inputs, **kwargs):
         model = cls(config, *inputs, **kwargs)
         return model
 

patched_padding_bert.py

@@ -1,39 +0,0 @@
-"""Source https://github.com/Dao-AILab/flash-attention/blob/87a1277653fc55cd615f5341255e00c69d5c00a1/flash_attn/bert_padding.py
-
-We replace the gather in `IndexFirstAxis.forward` with an indexing operation
-"""
-import torch
-from einops import rearrange, repeat
-
-class IndexFirstAxis(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, input, indices, indexing=False):
-        ctx.save_for_backward(indices)
-        assert input.ndim >= 2
-        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
-        # second_dim = other_shape.numel()
-        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
-        # return input[indices]
-        #return torch.gather(
-        #    rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim)
-        #).reshape(-1, *other_shape)
-        return input[indices]
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        (indices,) = ctx.saved_tensors
-        assert grad_output.ndim >= 2
-        other_shape = grad_output.shape[1:]
-        grad_output = rearrange(grad_output, "b ... -> b (...)")
-        grad_input = torch.zeros(
-            [ctx.first_axis_dim, grad_output.shape[1]],
-            device=grad_output.device,
-            dtype=grad_output.dtype,
-        )
-        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
-        # grad_input[indices] = grad_output
-        grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output)
-        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
-
-
-index_first_axis = IndexFirstAxis.apply