Tu2003716
/

COCOM_disabled_flash_attn

Model card Files Files and versions Community

Tu2003716 commited on about 1 month ago

Commit

e203c18

·

verified ·

1 Parent(s): e1ddb3d

Upload modeling_cocom.py

Files changed (1) hide show

modeling_cocom.py +1 -1

modeling_cocom.py CHANGED Viewed

@@ -226,7 +226,7 @@ class COCOM(PreTrainedModel):
         self.sep = cfg.sep
         self.compr_rate = cfg.compr_rate
         self.local_rank = os.getenv('LOCAL_RANK', '0')
-        for layer in self.decoder.layers:
             layer.attention.self = CustomFlashAttention(
                 embed_dim=cfg.hidden_size,
                 num_heads=cfg.num_attention_heads,

         self.sep = cfg.sep
         self.compr_rate = cfg.compr_rate
         self.local_rank = os.getenv('LOCAL_RANK', '0')
+        for layer in self.decoder.encoder.layer:
             layer.attention.self = CustomFlashAttention(
                 embed_dim=cfg.hidden_size,
                 num_heads=cfg.num_attention_heads,