Update modeling_binaryllm.py

modeling_binaryllm.py

@@ -11,6 +11,13 @@ from transformers.modeling_outputs import CausalLMOutput
 
 from .configuration_binaryllm import BinaryLLMConfig
 
+try:
+    import flash_attn_v100_cuda
+    _FLASH_V100_AVAILABLE = True
+except Exception:
+    flash_attn_v100_cuda = None
+    _FLASH_V100_AVAILABLE = False
+
 
 class PositionalEncoding(nn.Module):
     """
@@ -23,17 +30,17 @@ class PositionalEncoding(nn.Module):
         pe = torch.zeros(max_len, d_model, dtype=torch.float32)
         position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
         div_term = torch.exp(
-            torch.arange(0, d_model, 2, dtype=torch.float32) * (-torch.log(torch.tensor(10000.0)) / d_model)
+            torch.arange(0, d_model, 2, dtype=torch.float32)
+            * (-torch.log(torch.tensor(10000.0)) / d_model)
         )
         pe[:, 0::2] = torch.sin(position * div_term)
         pe[:, 1::2] = torch.cos(position * div_term)
-        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
+        pe = pe.unsqueeze(0)
         self.register_buffer("pe", pe, persistent=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         t = x.size(1)
-        pe = self.pe[:, :t, :]
-        pe = pe.to(device=x.device, dtype=x.dtype)
+        pe = self.pe[:, :t, :].to(device=x.device, dtype=x.dtype)
         return x + pe
 
 
@@ -48,6 +55,350 @@ class _InnerCfg:
     dropout: float
     layernorm_dim: Optional[int] = None
     head_dim: Optional[int] = None
+    attn_backend: str = "auto"
+
+
+class FlashSelfAttentionPortable(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        causal: bool = True,
+        backend: str = "auto",
+    ) -> None:
+        super().__init__()
+
+        if embed_dim % num_heads != 0:
+            raise ValueError(
+                f"embed_dim ({embed_dim}) doit être divisible par num_heads ({num_heads})"
+            )
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.dropout = float(dropout)
+        self.causal = bool(causal)
+        self.backend = str(backend)
+        self.softmax_scale = self.head_dim ** -0.5
+
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+
+    def _shape_qkv(
+        self,
+        x: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.dtype]:
+        bsz, seqlen, _ = x.shape
+        residual_dtype = x.dtype
+
+        proj_dtype = self.q_proj.weight.dtype
+        if x.dtype != proj_dtype:
+            x = x.to(proj_dtype)
+
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+
+        q = q.view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+        k = k.view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+        v = v.view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+        return q, k, v, residual_dtype
+
+    def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
+        bsz, nheads, seqlen, head_dim = x.shape
+        return x.transpose(1, 2).contiguous().view(bsz, seqlen, nheads * head_dim)
+
+    def _can_use_v100_kernel(self, q: torch.Tensor, padding_mask: Optional[torch.Tensor]) -> bool:
+        if not _FLASH_V100_AVAILABLE:
+            return False
+
+        if not q.is_cuda:
+            return False
+
+        if padding_mask is not None and bool(padding_mask.any().item()):
+            return False
+
+        cc = torch.cuda.get_device_capability(q.device)
+        if cc != (7, 0):
+            return False
+
+        hd = q.size(-1)
+        if hd % 2 != 0:
+            return False
+        if hd % 8 != 0:
+            return False
+        if hd > 256:
+            return False
+
+        return True
+
+    def _flash_attn_v100(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+    ) -> torch.Tensor:
+        if q.dtype != torch.float16:
+            q = q.to(torch.float16)
+        if k.dtype != torch.float16:
+            k = k.to(torch.float16)
+        if v.dtype != torch.float16:
+            v = v.to(torch.float16)
+
+        result = flash_attn_v100_cuda.fwd(
+            q,
+            k,
+            v,
+            None,
+            None,
+            0.0,
+            self.softmax_scale,
+            self.causal,
+            -1,
+            -1,
+            0.0,
+            False,
+            None,
+        )
+
+        out = result[0]
+        return out
+
+    def _sdpa_attn(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        bsz, nheads, tq, _ = q.shape
+        tk = k.size(-2)
+
+        attn_mask = None
+        if padding_mask is not None:
+            key_mask = padding_mask[:, None, None, :].to(device=q.device, dtype=torch.bool)
+            key_mask = key_mask.expand(bsz, nheads, tq, tk)
+            attn_mask = ~key_mask
+
+        dropout_p = self.dropout if self.training else 0.0
+
+        with torch.backends.cuda.sdp_kernel(
+            enable_flash=True,
+            enable_mem_efficient=True,
+            enable_math=True,
+        ):
+            out = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=attn_mask,
+                dropout_p=dropout_p,
+                is_causal=self.causal if attn_mask is None else False,
+                scale=self.softmax_scale,
+            )
+
+        return out
+
+    def _eager_attn(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * self.softmax_scale
+
+        if self.causal:
+            tq = q.size(-2)
+            tk = k.size(-2)
+            causal_mask = torch.triu(
+                torch.ones(tq, tk, device=scores.device, dtype=torch.bool),
+                diagonal=1,
+            )
+            scores = scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
+
+        if padding_mask is not None:
+            key_mask = padding_mask[:, None, None, :].to(device=scores.device, dtype=torch.bool)
+            scores = scores.masked_fill(key_mask, float("-inf"))
+
+        probs = torch.softmax(scores, dim=-1)
+
+        if self.training and self.dropout > 0.0:
+            probs = F.dropout(probs, p=self.dropout)
+
+        out = torch.matmul(probs, v.float())
+        return out.to(q.dtype)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        q, k, v, residual_dtype = self._shape_qkv(x)
+
+        if padding_mask is not None:
+            padding_mask = padding_mask.to(device=x.device, dtype=torch.bool)
+
+        backend = self.backend
+
+        if backend == "v100":
+            if not self._can_use_v100_kernel(q, padding_mask):
+                raise RuntimeError(
+                    "backend='v100' demandé mais indisponible "
+                    "(flash_attn_v100_cuda absent, GPU non sm70/V100, padding présent, "
+                    "ou head_dim incompatible)."
+                )
+            out = self._flash_attn_v100(q, k, v)
+
+        elif backend == "sdpa":
+            out = self._sdpa_attn(q, k, v, padding_mask=padding_mask)
+
+        elif backend == "eager":
+            out = self._eager_attn(q, k, v, padding_mask=padding_mask)
+
+        elif backend == "auto":
+            if self._can_use_v100_kernel(q, padding_mask):
+                out = self._flash_attn_v100(q, k, v)
+            else:
+                out = self._sdpa_attn(q, k, v, padding_mask=padding_mask)
+
+        else:
+            raise ValueError(f"backend d'attention non supporté: {backend}")
+
+        out = self._merge_heads(out)
+
+        out_proj_dtype = self.out_proj.weight.dtype
+        if out.dtype != out_proj_dtype:
+            out = out.to(out_proj_dtype)
+
+        out = self.out_proj(out)
+
+        if out.dtype != residual_dtype:
+            out = out.to(residual_dtype)
+
+        return out
+
+
+class FlashTransformerEncoderLayerPortable(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int,
+        dropout: float = 0.1,
+        activation: str = "gelu",
+        batch_first: bool = True,
+        attn_backend: str = "auto",
+    ) -> None:
+        super().__init__()
+
+        if not batch_first:
+            raise ValueError("Cette implémentation supporte batch_first=True uniquement.")
+
+        self.self_attn = FlashSelfAttentionPortable(
+            embed_dim=d_model,
+            num_heads=nhead,
+            dropout=dropout,
+            causal=True,
+            backend=attn_backend,
+        )
+
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+
+        self.dropout = nn.Dropout(dropout)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        if activation == "gelu":
+            self.activation = F.gelu
+        elif activation == "relu":
+            self.activation = F.relu
+        else:
+            raise ValueError(f"activation non supportée: {activation}")
+
+    def _sa_block(
+        self,
+        x: torch.Tensor,
+        src_key_padding_mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        x = self.self_attn(x, padding_mask=src_key_padding_mask)
+        x = self.dropout1(x)
+        return x
+
+    def _ff_block(self, x: torch.Tensor) -> torch.Tensor:
+        ff_dtype = self.linear1.weight.dtype
+        x_ff = x if x.dtype == ff_dtype else x.to(ff_dtype)
+
+        x_ff = self.linear1(x_ff)
+        x_ff = self.activation(x_ff)
+        x_ff = self.dropout(x_ff)
+        x_ff = self.linear2(x_ff)
+        x_ff = self.dropout2(x_ff)
+
+        if x_ff.dtype != x.dtype:
+            x_ff = x_ff.to(x.dtype)
+
+        return x_ff
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        x = src
+        x = self.norm1(x + self._sa_block(x, src_key_padding_mask))
+        x = self.norm2(x + self._ff_block(x))
+        return x
+
+
+class FlashTransformerEncoderPortable(nn.Module):
+    def __init__(
+        self,
+        encoder_layer: FlashTransformerEncoderLayerPortable,
+        num_layers: int,
+        attn_backend: str = "auto",
+    ) -> None:
+        super().__init__()
+
+        d_model = encoder_layer.norm1.normalized_shape[0]
+        nhead = encoder_layer.self_attn.num_heads
+        dim_feedforward = encoder_layer.linear1.out_features
+        dropout = encoder_layer.dropout.p
+
+        self.layers = nn.ModuleList(
+            [
+                FlashTransformerEncoderLayerPortable(
+                    d_model=d_model,
+                    nhead=nhead,
+                    dim_feedforward=dim_feedforward,
+                    dropout=dropout,
+                    activation="gelu",
+                    batch_first=True,
+                    attn_backend=attn_backend,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        x = src
+        for layer in self.layers:
+            x = layer(x, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
+        return x
 
 
 class TinyTransformerLM(nn.Module):
@@ -59,15 +410,20 @@ class TinyTransformerLM(nn.Module):
         self.tok_embed = nn.Embedding(vocab_size, cfg.embed_dim)
         self.pos_encoding = PositionalEncoding(cfg.embed_dim, cfg.block_size)
 
-        encoder_layer = nn.TransformerEncoderLayer(
+        encoder_layer = FlashTransformerEncoderLayerPortable(
             d_model=cfg.embed_dim,
             nhead=cfg.num_heads,
             dim_feedforward=cfg.ff_hidden_dim,
             dropout=cfg.dropout,
             activation="gelu",
             batch_first=True,
+            attn_backend=cfg.attn_backend,
+        )
+        self.encoder = FlashTransformerEncoderPortable(
+            encoder_layer,
+            num_layers=cfg.num_layers,
+            attn_backend=cfg.attn_backend,
         )
-        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=cfg.num_layers)
 
         ln_dim = cfg.layernorm_dim or cfg.embed_dim
         head_dim = cfg.head_dim or ln_dim
@@ -84,7 +440,6 @@ class TinyTransformerLM(nn.Module):
 
         self.head = nn.Linear(head_dim, vocab_size, bias=False)
 
-        # weight tying seulement si parfait alignement
         if self.pre_ln_proj is None and self.head_pre is None and head_dim == cfg.embed_dim:
             self.head.weight = self.tok_embed.weight
 
@@ -104,13 +459,26 @@ class TinyTransformerLM(nn.Module):
         x = self.encoder(x, mask=attn_mask, src_key_padding_mask=padding_mask)
 
         if self.pre_ln_proj is not None:
+            proj_dtype = self.pre_ln_proj.weight.dtype
+            if x.dtype != proj_dtype:
+                x = x.to(proj_dtype)
             x = self.pre_ln_proj(x)
 
+        ln_dtype = self.ln.weight.dtype
+        if x.dtype != ln_dtype:
+            x = x.to(ln_dtype)
         x = self.ln(x)
 
         if self.head_pre is not None:
+            head_pre_dtype = self.head_pre.weight.dtype
+            if x.dtype != head_pre_dtype:
+                x = x.to(head_pre_dtype)
             x = self.head_pre(x)
 
+        head_dtype = self.head.weight.dtype
+        if x.dtype != head_dtype:
+            x = x.to(head_dtype)
+
         return self.head(x)
 
 
@@ -121,6 +489,8 @@ class BinaryLLMForCausalLM(PreTrainedModel):
     def __init__(self, config: BinaryLLMConfig):
         super().__init__(config)
 
+        attn_backend = getattr(config, "attn_backend", "auto")
+
         inner = _InnerCfg(
             block_size=int(config.max_position_embeddings),
             embed_dim=int(config.hidden_size),
@@ -131,11 +501,35 @@ class BinaryLLMForCausalLM(PreTrainedModel):
             dropout=float(getattr(config, "dropout", 0.0)),
             layernorm_dim=None,
             head_dim=None,
+            attn_backend=str(attn_backend),
         )
         self.model = TinyTransformerLM(inner)
 
         self.post_init()
 
+    def get_input_embeddings(self) -> nn.Module:
+        return self.model.tok_embed
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.model.tok_embed = value
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.model.head
+
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+        self.model.head = new_embeddings
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+        }
+
     def forward(
         self,
         input_ids: torch.LongTensor,
@@ -145,7 +539,7 @@ class BinaryLLMForCausalLM(PreTrainedModel):
     ) -> CausalLMOutput:
         padding_mask = None
         if attention_mask is not None:
-            padding_mask = ~attention_mask.to(torch.bool)  # True = ignore
+            padding_mask = ~attention_mask.to(torch.bool)
 
         logits = self.model(input_ids, padding_mask=padding_mask)