Fix: all_tied_weights_keys

modeling_ezfgraphic.py

@@ -1,91 +1,93 @@
-from .configuration_ezfgraphic import EZFGraphicConfig
-
-import torch, torch.nn as nn, torch.nn.functional as F
-from transformers import PreTrainedModel
-
-class RMSNorm(nn.Module):
-    def __init__(self, d, eps=1e-5):
-        super().__init__()
-        self.w = nn.Parameter(torch.ones(d)); self.eps = eps
-    def forward(self, x):
-        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.w
-
-class Attention(nn.Module):
-    def __init__(self, cfg):
-        super().__init__()
-        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
-        self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
-        self.drop = nn.Dropout(cfg.dropout)
-        self.register_buffer("mask", torch.tril(torch.ones(512, 512)).bool())
-    def forward(self, x):
-        B, T, C = x.shape
-        q, k, v = self.qkv(x).chunk(3, dim=-1)
-        att = (q @ k.transpose(-2, -1)) / (C ** 0.5)
-        att = att.masked_fill(~self.mask[:T, :T], float("-inf"))
-        att = self.drop(F.softmax(att, dim=-1))
-        return self.proj(att @ v)
-
-class FeedForward(nn.Module):
-    def __init__(self, cfg):
-        super().__init__()
-        self.fc1 = nn.Linear(cfg.d_model, 2 * cfg.d_model)
-        self.fc2 = nn.Linear(2 * cfg.d_model, cfg.d_model)
-        self.drop = nn.Dropout(cfg.dropout)
-    def forward(self, x):
-        return self.drop(self.fc2(F.relu(self.fc1(x))))
-
-class Block(nn.Module):
-    def __init__(self, cfg):
-        super().__init__()
-        self.ln1 = RMSNorm(cfg.d_model); self.attn = Attention(cfg)
-        self.ln2 = RMSNorm(cfg.d_model); self.ffn = FeedForward(cfg)
-    def forward(self, x):
-        x = x + self.attn(self.ln1(x))
-        x = x + self.ffn(self.ln2(x))
-        return x
-
-class EZFGraphic(PreTrainedModel):
-    config_class = EZFGraphicConfig
-    def __init__(self, config):
-        super().__init__(config)
-        cfg = config
-        self.cfg = cfg
-        self.text_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
-        self.pos_emb_text = nn.Embedding(cfg.ctx, cfg.d_model)
-        n_patches = (cfg.image_size // cfg.patch_size) ** 2
-        self.n_patches = n_patches
-        self.pos_emb_visual = nn.Embedding(n_patches, cfg.d_model)
-        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
-        self.ln_f = RMSNorm(cfg.d_model)
-        self.patch_pixels = cfg.patch_size * cfg.patch_size * 3
-        self.pixel_decoder = nn.Sequential(
-            nn.Linear(cfg.d_model, 512),
-            nn.ReLU(),
-            nn.Linear(512, 1024),
-            nn.ReLU(),
-            nn.Linear(1024, self.patch_pixels),
-        )
-    def encode_text(self, text_ids):
-        B, T = text_ids.shape
-        pos = torch.arange(T, device=text_ids.device)
-        x = self.text_emb(text_ids) + self.pos_emb_text(pos)
-        for b in self.blocks:
-            x = b(x)
-        return self.ln_f(x)
-    def decode_pixels(self, x):
-        B, T, _ = x.shape
-        patches = self.pixel_decoder(x)
-        ps = self.cfg.patch_size
-        h = w = self.cfg.image_size // ps
-        patches = patches.view(B, h, w, ps, ps, 3)
-        pixels = patches.permute(0, 5, 1, 3, 2, 4).reshape(B, 3, self.cfg.image_size, self.cfg.image_size)
-        return pixels
-    def forward(self, text_ids):
-        x = self.encode_text(text_ids)
-        last = x[:, -1:, :]
-        visual_input = last.expand(-1, self.n_patches, -1).contiguous()
-        visual_input = visual_input + self.pos_emb_visual(
-            torch.arange(self.n_patches, device=text_ids.device)
-        ).unsqueeze(0)
-        pixels = self.decode_pixels(visual_input)
-        return pixels
+from .configuration_ezfgraphic import EZFGraphicConfig
+
+import torch, torch.nn as nn, torch.nn.functional as F
+from transformers import PreTrainedModel
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-5):
+        super().__init__()
+        self.w = nn.Parameter(torch.ones(d)); self.eps = eps
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.w
+
+class Attention(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
+        self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self.drop = nn.Dropout(cfg.dropout)
+        self.register_buffer("mask", torch.tril(torch.ones(512, 512)).bool())
+    def forward(self, x):
+        B, T, C = x.shape
+        q, k, v = self.qkv(x).chunk(3, dim=-1)
+        att = (q @ k.transpose(-2, -1)) / (C ** 0.5)
+        att = att.masked_fill(~self.mask[:T, :T], float("-inf"))
+        att = self.drop(F.softmax(att, dim=-1))
+        return self.proj(att @ v)
+
+class FeedForward(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.fc1 = nn.Linear(cfg.d_model, 2 * cfg.d_model)
+        self.fc2 = nn.Linear(2 * cfg.d_model, cfg.d_model)
+        self.drop = nn.Dropout(cfg.dropout)
+    def forward(self, x):
+        return self.drop(self.fc2(F.relu(self.fc1(x))))
+
+class Block(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.ln1 = RMSNorm(cfg.d_model); self.attn = Attention(cfg)
+        self.ln2 = RMSNorm(cfg.d_model); self.ffn = FeedForward(cfg)
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.ffn(self.ln2(x))
+        return x
+
+class EZFGraphic(PreTrainedModel):
+    config_class = EZFGraphicConfig
+    _tied_weights_keys = []
+    all_tied_weights_keys = {}
+    def __init__(self, config):
+        super().__init__(config)
+        cfg = config
+        self.cfg = cfg
+        self.text_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.pos_emb_text = nn.Embedding(cfg.ctx, cfg.d_model)
+        n_patches = (cfg.image_size // cfg.patch_size) ** 2
+        self.n_patches = n_patches
+        self.pos_emb_visual = nn.Embedding(n_patches, cfg.d_model)
+        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
+        self.ln_f = RMSNorm(cfg.d_model)
+        self.patch_pixels = cfg.patch_size * cfg.patch_size * 3
+        self.pixel_decoder = nn.Sequential(
+            nn.Linear(cfg.d_model, 512),
+            nn.ReLU(),
+            nn.Linear(512, 1024),
+            nn.ReLU(),
+            nn.Linear(1024, self.patch_pixels),
+        )
+    def encode_text(self, text_ids):
+        B, T = text_ids.shape
+        pos = torch.arange(T, device=text_ids.device)
+        x = self.text_emb(text_ids) + self.pos_emb_text(pos)
+        for b in self.blocks:
+            x = b(x)
+        return self.ln_f(x)
+    def decode_pixels(self, x):
+        B, T, _ = x.shape
+        patches = self.pixel_decoder(x)
+        ps = self.cfg.patch_size
+        h = w = self.cfg.image_size // ps
+        patches = patches.view(B, h, w, ps, ps, 3)
+        pixels = patches.permute(0, 5, 1, 3, 2, 4).reshape(B, 3, self.cfg.image_size, self.cfg.image_size)
+        return pixels
+    def forward(self, text_ids):
+        x = self.encode_text(text_ids)
+        last = x[:, -1:, :]
+        visual_input = last.expand(-1, self.n_patches, -1).contiguous()
+        visual_input = visual_input + self.pos_emb_visual(
+            torch.arange(self.n_patches, device=text_ids.device)
+        ).unsqueeze(0)
+        pixels = self.decode_pixels(visual_input)
+        return pixels