feat: CP support for mamba layer (#164)

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <[email protected]>

* debug

Signed-off-by: Mehant Kammakomati <[email protected]>

* debug

Signed-off-by: Mehant Kammakomati <[email protected]>

* fix: remove print stmts

Signed-off-by: Mehant Kammakomati <[email protected]>

* docs: add docs

Signed-off-by: Mehant Kammakomati <[email protected]>

* docs: add docs

Signed-off-by: Mehant Kammakomati <[email protected]>

* nit: lint and fmt

Signed-off-by: Mehant Kammakomati <[email protected]>

* fix: dp cp loss

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add unit test

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add unit test

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add unit test

Signed-off-by: Mehant Kammakomati <[email protected]>

* feat: add unit test

Signed-off-by: Mehant Kammakomati <[email protected]>

* nit: remove ds_store files

Signed-off-by: Mehant Kammakomati <[email protected]>

---------

Signed-off-by: Mehant Kammakomati <[email protected]>

Author: kmehant

.github/workflows/build-and-publish.yml

@@ -17,6 +17,7 @@ jobs:
           - "attention-and-distributed-packing"
           - "accelerated-moe"
           - "online-data-mixing"
+          - "mamba-cp"
 
     permissions:
       id-token: write  # IMPORTANT: this permission is mandatory for trusted publishing

.github/workflows/format.yml

@@ -32,6 +32,7 @@ jobs:
           - "attention-and-distributed-packing"
           - "accelerated-moe"
           - "online-data-mixing"
+          - "mamba-cp"
 
     steps:
       - name: Delete huge unnecessary tools folder

.gitignore

@@ -1,4 +1,5 @@
 __pycache__
 *.tar.gz
 *.tox
-*.pytest_cache
+*.pytest_cache
+**/.DS_Store

plugins/framework/src/fms_acceleration/constants.py

@@ -21,4 +21,4 @@
 # and activated.
 # - hence the plugins that have model loaders should be on top of this list
 
-PLUGINS = ["peft", "foak", "aadp", "moe", "odm"]
+PLUGINS = ["peft", "foak", "aadp", "moe", "odm", "mcp"]

plugins/mamba-cp/.isort.cfg

@@ -0,0 +1,10 @@
+[settings]
+profile=black
+from_first=true
+import_heading_future=Future
+import_heading_stdlib=Standard
+import_heading_thirdparty=Third Party
+import_heading_firstparty=First Party
+import_heading_localfolder=Local
+known_firstparty=
+known_localfolder=tuning