Upload BD3LM

modeling_bd3lm.py

@@ -299,7 +299,7 @@ class DDiTBlock(nn.Module):
       return bias_dropout_add_scale_fused_inference
   
 
-  def get_qkv(self, x, rotary_cos_sin, save_kv=False):
+  def get_qkv(self, x, rotary_cos_sin, store_kv=False):
     # compute qkv (potentially use cache)
     if self.kv_cache is not None:
       block_len = x.shape[1] - self.kv_cache.shape[1]
@@ -308,8 +308,8 @@ class DDiTBlock(nn.Module):
     else:
       qkv = self.attn_qkv(x)
     
-    # save kv cache in a sliding window (can't exceed context len)
-    if save_kv:
+    # store kv cache in a sliding window (can't exceed context len)
+    if store_kv:
       if self.kv_cache is not None:
         cache_len = min(x.shape[1], self.n - block_len)
         self.kv_cache = qkv[:, -cache_len:]
@@ -347,7 +347,8 @@ class DDiTBlock(nn.Module):
     x = einops.rearrange(x, 'b s h d -> b s (h d)')
     return x
 
-  def forward(self, x, rotary_cos_sin, c, cross_attn_mask=None, save_kv=False):
+  def forward(self, x, rotary_cos_sin, c, cross_attn_mask=None,
+              sample_mode=False, store_kv=False):
     bias_dropout_scale_fn = self._get_bias_dropout_scale()
 
     (shift_msa, scale_msa, gate_msa, shift_mlp,
@@ -358,12 +359,12 @@ class DDiTBlock(nn.Module):
     x = modulate_fused(self.norm1(x), shift_msa, scale_msa)
 
     # get qkvs
-    if cross_attn_mask is not None and not save_kv:
+    if cross_attn_mask is not None and not sample_mode:
       qkv_x = self.get_qkv(x[:,:self.n], rotary_cos_sin)
       qkv_x0 = self.get_qkv(x[:,self.n:], rotary_cos_sin)
       qkv = torch.cat((qkv_x, qkv_x0), dim=1)
     else:
-      qkv = self.get_qkv(x, rotary_cos_sin, save_kv=save_kv)
+      qkv = self.get_qkv(x, rotary_cos_sin, store_kv=store_kv)
 
     if cross_attn_mask is None and self.attn_backend == 'flash_attn':
       x = regular_attention_multi_headed(qkv)
@@ -470,9 +471,8 @@ class DITBackbone(nn.Module):
     x0_attn_mask = torch.cat((torch.zeros_like(self_attn_mask), x0_attn_mask), dim=1)
     self.cross_attn_mask = torch.cat((cross_attn_mask, x0_attn_mask), dim=0)
 
-  def forward(self, indices, sigma, disable_cross_attn=False,
-              output_hidden_states=False, save_kv=False):
-    cross_attn = self.cross_attn and not disable_cross_attn
+  def forward(self, indices, sigma, sample_mode=False,
+             store_kv=False, output_hidden_states=False):
     if not self.config.time_conditioning:
       sigma = torch.zeros_like(sigma)
     all_hidden_states = []
@@ -480,11 +480,13 @@ class DITBackbone(nn.Module):
     if output_hidden_states:
       all_hidden_states.append(x)
     c = F.silu(self.sigma_map(sigma))
-    if cross_attn:
-      cross_attn_mask = self.cross_attn_mask.to(x.device)
-      if save_kv:
-        cross_attn_mask = cross_attn_mask[:x.shape[1], :x.shape[1]]
+    if self.cross_attn:
       rotary_cos_sin = self.rotary_emb(x[:, :self.n])
+      cross_attn_mask = self.cross_attn_mask.to(x.device)
+      # use block-causal mask only during sampling
+      if sample_mode:
+        cross_attn_mask = cross_attn_mask[
+          self.n:self.n+x.shape[1], self.n:self.n+x.shape[1]]
     else:
       cross_attn_mask = None
       rotary_cos_sin = self.rotary_emb(x)
@@ -495,11 +497,12 @@ class DITBackbone(nn.Module):
                            rotary_cos_sin,
                            c,
                            cross_attn_mask=cross_attn_mask,
-                           save_kv=save_kv)
+                           sample_mode=sample_mode,
+                           store_kv=store_kv)
         if output_hidden_states:
           all_hidden_states.append(x)
       logits = self.output_layer(x, c)
-    if cross_attn and not save_kv:
+    if self.cross_attn and not sample_mode:
       logits = logits[:, :self.n]
       all_hidden_states = [hidden_states[:, :self.n] for hidden_states in all_hidden_states]
     return logits, all_hidden_states
@@ -526,7 +529,8 @@ class BD3LM(transformers.PreTrainedModel):
       self,
       input_ids: torch.LongTensor = None,
       timesteps: torch.FloatTensor = None,
-      disable_cross_attn: typing.Optional[bool] = None,
+      sample_mode: typing.Optional[bool] = None,
+      store_kv: typing.Optional[bool] = None,
       output_hidden_states: typing.Optional[bool] = None,
       return_dict: typing.Optional[bool] = None,
   ) -> typing.Union[
@@ -545,8 +549,9 @@ class BD3LM(transformers.PreTrainedModel):
     logits, all_hidden_states = self.backbone(
       indices=input_ids,
       sigma=timesteps,
-      disable_cross_attn=disable_cross_attn,
-      output_hidden_states=output_hidden_states
+      sample_mode=sample_mode,
+      store_kv=store_kv,
+      output_hidden_states=output_hidden_states,
     )
     if return_dict:
       return modeling_outputs.MaskedLMOutput(