feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>

Author: kmehant

plugins/mamba-cp/src/fms_acceleration_mcp/framework_plugin_mcp.py

@@ -21,6 +21,7 @@
 from transformers import TrainingArguments
 import torch
 
+# Local
 from .utils import patch_mamba_layers_with_cp_head
 
 
@@ -41,6 +42,7 @@ def __init__(self, configurations: Dict[str, Dict]):
             key="training.mamba.cp.mamba_recompute",
             default=False,
         )
+
     # data_config file should be there
     @property
     def requires_augmentation(self):
@@ -52,7 +54,7 @@ def augmentation(
         train_args: TrainingArguments,
         modifiable_args: Tuple[LoraConfig],
     ):
-        if self._mamba_cp_degree != None:
+        if self._mamba_cp_degree is not None:
             rank = 0
             if torch.distributed.is_initialized():
                 rank = torch.distributed.get_node_local_rank()

plugins/mamba-cp/src/fms_acceleration_mcp/utils/__init__.py

@@ -14,4 +14,3 @@
 
 # Local
 from .utils import patch_mamba_layers_with_cp_head
-

plugins/mamba-cp/src/fms_acceleration_mcp/utils/utils.py

@@ -11,18 +11,17 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+# Standard
 from typing import Dict
 
+# Third Party
+from mamba_ssm.modules.mamba2_cp import Mamba2CP
+
 # pylint: disable=import-error
-from torch.distributed._tensor.device_mesh import DeviceMesh, init_device_mesh
-import torch
-from transformers.modeling_utils import is_fsdp_enabled, is_local_dist_rank_0
+from torch.distributed._tensor.device_mesh import init_device_mesh
 from tqdm import tqdm
-
-try:
-    from mamba_ssm.modules.mamba2_cp import Mamba2CP
-except ImportError:
-    ValueError("Mamba2CP is required to enable context parallelism for mamba layers")
+from transformers.modeling_utils import is_fsdp_enabled
+import torch
 
 key_ep = "cp"
 key_rep = "dp_shard"
@@ -42,9 +41,22 @@ def hf_config_ssm_config(hf_config) -> Dict:
 
 class Mamba2CPHF(Mamba2CP):
     def forward(
-        self, hidden_states, cache_params=None, cache_position=None, attention_mask=None, seq_idx=None, **kwargs
+        self,
+        hidden_states,
+        cache_params=None,
+        cache_position=None,
+        attention_mask=None,
+        seq_idx=None,
+        **kwargs,
     ):
-        return super().forward(u=hidden_states, seqlen=None, seq_idx=None, cu_seqlens=None, inference_params=None)
+        return super().forward(
+            u=hidden_states,
+            seqlen=None,
+            seq_idx=None,
+            cu_seqlens=None,
+            inference_params=None,
+        )
+
 
 def patch_mamba_layers_with_cp_head(
     model,
@@ -63,7 +75,7 @@ def patch_mamba_layers_with_cp_head(
 
     if cp_degree == 1:
         raise ValueError("CP degree can't be one")
-    elif rep_size == 1:
+    if rep_size == 1:
         device_mesh = init_device_mesh(
             "cuda",
             (cp_degree,),