Update API

Signed-off-by: Dashiell Stander <[email protected]>
EleutherAI · Oct 3, 2023 · 6189c0d · 6189c0d
1 parent 3f43f07
commit 6189c0d
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -344,7 +344,6 @@ def __init__(
                 neox_args,
                 self.attention_type,
                 self.num_attention_heads_per_partition,
-                mpu=mpu,
             )
         else:
             if self.use_flash_attention:

diff --git a/megatron/model/utils.py b/megatron/model/utils.py
@@ -238,7 +238,7 @@ def _set_use_cache(modules, value: bool):
     recursive_setattr(modules, "use_cache", value, assert_type=bool)
 
 
-def configure_sparse_attention(neox_args, attention_type, num_attention_heads, mpu):
+def configure_sparse_attention(neox_args, attention_type, num_attention_heads):
     from deepspeed.ops.sparse_attention import (
         SparseSelfAttention,
         VariableSparsityConfig,
@@ -337,7 +337,6 @@ def configure_sparse_attention(neox_args, attention_type, num_attention_heads, m
         sparsity_config=sparsity_config,
         max_seq_length=neox_args.seq_length,
         attn_mask_mode="add",
-        mpu=mpu,
     )