remove optional flag (NVIDIA-NeMo#10781)

* remove optional flag

* temp fix CI

* Apply isort and black reformatting

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

* add triton==2.2.0 to reqs

* undo reqs

* add triton==2.2.0 to reqs

* remove triton downgrade from dockerfile

* bump mamba version

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Apply isort and black reformatting

Signed-off-by: pre-commit-ci[bot] <pre-commit-ci[bot]@users.noreply.github.com>

---------

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>
Signed-off-by: pre-commit-ci[bot] <pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ali Taghibakhshi <ataghibakhsh@login-eos01.eos.clusters.nvidia.com>
Co-authored-by: JRD971000 <JRD971000@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

.github/workflows/cicd-main.yml

@@ -5151,10 +5151,10 @@ jobs:
         rm -rf tests/collections/llm/gpt_index_mappings
       IS_OPTIONAL: true
 
-  OPTIONAL_L2_NeMo_2_SSM_Pretraining:
+  L2_NeMo_2_SSM_Pretraining:
     needs: [cicd-test-container-setup]
     uses: ./.github/workflows/_test_template.yml
-    if: contains(fromJSON(needs.cicd-test-container-setup.outputs.test_to_run), 'OPTIONAL_L2_NeMo_2_SSM_Pretraining') || needs.cicd-test-container-setup.outputs.all == 'true'
+    if: contains(fromJSON(needs.cicd-test-container-setup.outputs.test_to_run), 'L2_NeMo_2_SSM_Pretraining') || needs.cicd-test-container-setup.outputs.all == 'true'
     with:
       RUNNER: self-hosted-azure
       SCRIPT: |
@@ -5164,12 +5164,11 @@ jobs:
         --max-steps 10 \
         --experiment-dir /tmp/nlp_megatron_mamba_nemo-ux-mamba_cicd_test_pretrain/${{ github.run_id }} \
         --data-path /home/TestData/nlp/megatron_mamba/toy_ssm_dataset/legal_pile_text_document
-      IS_OPTIONAL: true
 
-  OPTIONAL_L2_NeMo_2_SSM_Finetuning:
+  L2_NeMo_2_SSM_Finetuning:
     needs: [cicd-test-container-setup]
     uses: ./.github/workflows/_test_template.yml
-    if: contains(fromJSON(needs.cicd-test-container-setup.outputs.test_to_run), 'OPTIONAL_L2_NeMo_2_SSM_Finetuning') || needs.cicd-test-container-setup.outputs.all == 'true'
+    if: contains(fromJSON(needs.cicd-test-container-setup.outputs.test_to_run), 'L2_NeMo_2_SSM_Finetuning') || needs.cicd-test-container-setup.outputs.all == 'true'
     with:
       RUNNER: self-hosted-azure
       SCRIPT: |
@@ -5179,7 +5178,6 @@ jobs:
         --max-steps 10 \
         --experiment-dir /tmp/nlp_megatron_mamba_nemo-ux-mamba_cicd_test_sft/${{ github.run_id }} \
         --model-path /home/TestData/nlp/megatron_mamba/model_optim_rng.pt
-      IS_OPTIONAL: true
 
   L2_NeMo_2_HF_MODEL_IMPORT:
     needs: [cicd-test-container-setup]
@@ -5364,8 +5362,8 @@ jobs:
       - L2_NeMo_2_GPT_Pretraining_no_transformer_engine
       #- OPTIONAL_L2_NeMo_2_GPT_DDP_Param_Parity_check
       - L2_NeMo_2_HF_MODEL_IMPORT
-      #- OPTIONAL_L2_NeMo_2_SSM_Pretraining
-      #- OPTIONAL_L2_NeMo_2_SSM_Finetuning
+      - L2_NeMo_2_SSM_Pretraining
+      - L2_NeMo_2_SSM_Finetuning
       - L2_NeMo_2_T5_Pretraining
       - L2_NeMo_2_Mixtral_Pretraining
       - L2_PTQ_Llama2_INT8_SQ

Dockerfile.ci

@@ -32,17 +32,11 @@ EOF
 
 WORKDIR /workspace
 
-# Install Mamba
-ARG MAMBA_TAG=v2.0.3
+# Install Mamba Dependancy
 ARG CAUSAL_CONV_TAG=v1.2.2.post1
 
 RUN <<"EOF" bash -ex
 # Mamba dependancy installation
-git clone --depth 1 --branch ${MAMBA_TAG} https://github.com/state-spaces/mamba.git && \
-  cd mamba && \
-  python setup.py install && \
-  cd .. && \
-  rm -rf mamba
 
 git clone --depth 1 --branch ${CAUSAL_CONV_TAG} https://github.com/Dao-AILab/causal-conv1d && \
   cd causal-conv1d && \

requirements/requirements_nlp.txt

@@ -9,6 +9,7 @@ gdown
 h5py
 ijson
 jieba
+mamba-ssm==2.2.2
 markdown2
 matplotlib>=3.3.2
 #megatron_core>0.6.0 # add back once mcore on pypi is compatible again