File size: 2,206 Bytes
ed12f96 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | diff --git a/mindspeed/core/fusions/fused_rope.py b/mindspeed/core/fusions/fused_rope.py
index a6f02e07..70f7cb08 100644
--- a/mindspeed/core/fusions/fused_rope.py
+++ b/mindspeed/core/fusions/fused_rope.py
@@ -126,5 +126,6 @@ def apply_rotary_pos_emb(
freqs,
rotary_interleaved=config.rotary_interleaved,
multi_latent_attention=config.multi_latent_attention,
- mscale=mscale
+ mscale=mscale,
+ cp_group=cp_group
)
diff --git a/mindspeed/megatron_adaptor.py b/mindspeed/megatron_adaptor.py
index f14b231d..4615590e 100644
--- a/mindspeed/megatron_adaptor.py
+++ b/mindspeed/megatron_adaptor.py
@@ -57,6 +57,7 @@ def delete_lock_file():
def repatch(args):
MindSpeedFeaturesManager.remove_patches()
full_args = get_full_args()
+ args = vars(args)
for k, v in args.items():
setattr(full_args, k, v)
MindSpeedFeaturesManager.apply_features_pre_patches(full_args)
diff --git a/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py b/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
index ac4eabe5..78c5866d 100644
--- a/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -330,6 +330,8 @@ class DotProductAttention(torch.nn.Module):
inference_params: Any = None,
pad_between_seqs: Optional[bool] = None,
fp8_output: Optional[bool] = False,
+ local_cp_size=None,
+ cp_group=None,
) -> torch.Tensor:
"""
Dot Product Attention Layer.
@@ -659,7 +661,9 @@ class MindSpeedTEDotProductAttention(DotProductAttention):
) and not getattr(self.config, 'is_llava', False):
self.config.sparse_mode = 2
attention_mask = get_attention_mask(self.config)
-
+ attention_mask = torch.triu(
+ torch.ones((2048, 2048),
+ device=query.device, dtype=torch.bool), diagonal=1)
packed_seq_kwargs = (
{key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
if packed_seq_params is not None
|