Upload BD3LM

modeling_bd3lm.py

@@ -16,17 +16,15 @@ try:
   FLEX_ATTN_AVAILABLE = True
 except:
   FLEX_ATTN_AVAILABLE = False
-# Flags required to enable jit fusion kernels
-try:
-  torch._C._jit_set_profiling_mode(False)
-  torch._C._jit_set_profiling_executor(False)
-  torch._C._jit_override_can_fuse_on_cpu(True)
-  torch._C._jit_override_fcan_fuse_on_gpu(True)
-except:
-  pass
 
 from .configuration_bd3lm import BD3LMConfig
 
+# Flags required to enable jit fusion kernels
+torch._C._jit_set_profiling_mode(False)
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_override_can_fuse_on_cpu(True)
+torch._C._jit_override_can_fuse_on_gpu(True)
+
 def block_diff_mask(b, h, q_idx, kv_idx, block_size=None, n=None):
   """
   Constructs the specialized block diffusion attention mask for training
@@ -77,7 +75,6 @@ def block_diff_mask(b, h, q_idx, kv_idx, block_size=None, n=None):
 def fused_flex_attention(q, k, v, mask=None):
     return flex_attention(q, k, v, block_mask=mask)
 
-
 def bias_dropout_add_scale(
     x: torch.Tensor,
     bias: typing.Optional[torch.Tensor],
@@ -102,6 +99,7 @@ def get_bias_dropout_add_scale(training):
 
   return _bias_dropout_add
 
+
 # function overload
 def modulate(x: torch.Tensor,
              shift: torch.Tensor,
@@ -299,7 +297,7 @@ def regular_attention_multi_headed(qkv):
 
 class DDiTBlock(nn.Module):
   def __init__(self, n, block_size, dim, n_heads, cond_dim, mlp_ratio=4,
-               dropout=0.1, max_seqlen=1024, attn_backend='flash_attn'):
+               dropout=0.1, attn_backend='sdpa'):
     super().__init__()
     self.n = n
     self.block_size = block_size
@@ -394,11 +392,9 @@ class DDiTBlock(nn.Module):
     else:
       qkv = self.get_qkv(x, rotary_cos_sin, store_kv=store_kv)
 
-    if mask is None and self.attn_backend == 'flash_attn':
-      x = regular_attention_multi_headed(qkv)
-    elif self.attn_backend == 'flex' and FLEX_ATTN_AVAILABLE:
+    if self.attn_backend == 'flex' and FLEX_ATTN_AVAILABLE:
       x = self.cross_attn_flex(qkv, mask=mask)
-    elif self.attn_backend == 'sdpa':
+    elif self.attn_backend == 'sdpa' or not FLEX_ATTN_AVAILABLE:
       x = self.cross_attn(x, qkv, mask=mask)
     else:
       raise ValueError('Unknown attention backend')
@@ -500,7 +496,7 @@ class DITBackbone(nn.Module):
       self.mask = create_block_mask(
         partial(block_diff_mask, block_size=block_size, n=seqlen),
         B=None, H=None, Q_LEN=seqlen*2, KV_LEN=seqlen*2)
-    elif attn_backend == 'sdpa':
+    elif attn_backend == 'sdpa' or not FLEX_ATTN_AVAILABLE:
       self.mask = block_diff_mask(
         b=None, h=None, q_idx=torch.arange(seqlen*2)[:, None], 
         kv_idx=torch.arange(seqlen*2)[None, :], block_size=block_size, n=seqlen)