diff --git a/.coveragerc b/.coveragerc
deleted file mode 100644
index 0e37a5470..000000000
--- a/.coveragerc
+++ /dev/null
@@ -1,6 +0,0 @@
-[html]
-directory = coverage
-
-[run]
-data_file = .coverage_$LOCAL_RANK
-relative_files = true
diff --git a/.flake8 b/.flake8
deleted file mode 100644
index 1e35e0c49..000000000
--- a/.flake8
+++ /dev/null
@@ -1,4 +0,0 @@
-[flake8]
-max-line-length = 100
-extend-ignore = E203,E501,F401,E402,E714
-per-file-ignores = __init__.py:F401
\ No newline at end of file
diff --git a/.github/ISSUE_TEMPLATE/bug.md b/.github/ISSUE_TEMPLATE/bug.md
deleted file mode 100644
index b639acd3c..000000000
--- a/.github/ISSUE_TEMPLATE/bug.md
+++ /dev/null
@@ -1,32 +0,0 @@
----
-name: BUG
-about: Report a bug that needs attention
-title: "[BUG]"
-labels: ''
-assignees: ''
-
----
-
-**Describe the bug**
-A clear and concise description of what the bug is.
-
-**To Reproduce**
-Steps to reproduce the behavior. The easier it is to reproduce the faster it will get maintainer attention.
-
-**Expected behavior**
-A clear and concise description of what you expected to happen.
-
-**Stack trace/logs**
-If applicable, add the stack trace or logs from the time of the error.
-
-**Environment (please complete the following information):**
- - Megatron-LM commit ID
- - PyTorch version
- - CUDA version
- - NCCL version
-
-**Proposed fix**
-If you have a proposal for how to fix the issue state it here or link to a PR.
-
-**Additional context**
-Add any other context about the problem here.
diff --git a/.github/ISSUE_TEMPLATE/enhancement.md b/.github/ISSUE_TEMPLATE/enhancement.md
deleted file mode 100644
index 076f7195b..000000000
--- a/.github/ISSUE_TEMPLATE/enhancement.md
+++ /dev/null
@@ -1,23 +0,0 @@
----
-name: ENHANCEMENT
-about: Suggest an idea to improve this project
-title: "[ENHANCEMENT]"
-labels: ''
-assignees: ''
-
----
-
-**Is your feature request related to a problem? Please describe.**
-A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
-
-**Describe the solution you'd like**
-A clear and concise description of what you want to happen.
-
-**Describe alternatives you've considered**
-A clear and concise description of any alternative solutions or features you've considered.
-
-**Proposed implementation**
-If you have a proposed implementation for the feature state it here or link to a PR.
-
-**Additional context**
-Add any other context or screenshots about the feature request here.
diff --git a/.github/ISSUE_TEMPLATE/question.md b/.github/ISSUE_TEMPLATE/question.md
deleted file mode 100644
index b3d89a0ac..000000000
--- a/.github/ISSUE_TEMPLATE/question.md
+++ /dev/null
@@ -1,12 +0,0 @@
----
-name: QUESTION
-about: Ask a question about Megatron-LM that is not a bug, regression or enhancement
-  request
-title: "[QUESTION]"
-labels: ''
-assignees: ''
-
----
-
-**Your question**
-Ask a clear and concise question about Megatron-LM.
diff --git a/.github/ISSUE_TEMPLATE/regression.md b/.github/ISSUE_TEMPLATE/regression.md
deleted file mode 100644
index 10078d23a..000000000
--- a/.github/ISSUE_TEMPLATE/regression.md
+++ /dev/null
@@ -1,39 +0,0 @@
----
-name: REGRESSION
-about: Report a regression in speed or accuracy due to a Megatron-LM update
-title: "[REGRESSION]"
-labels: ''
-assignees: ''
-
----
-
-**Describe the regression**
-A clear and concise description of what the regression is.
-
-**To Reproduce**
-Steps to reproduce the behavior. The easier it is to reproduce the faster it will get maintainer attention.
-
-**Previous performance**
-What speed or accuracy did you previously see.
-
-**New performance**
-What speed or accuracy do you see after the update.
-
-**Stack trace/logs**
-If applicable, add the stack trace or logs related to the regression.
-
-**Environment (please complete the following information):**
- - Previous Megatron-LM commit ID
- - New Megatron-LM commit ID
- - Previous PyTorch version
- - New PyTorch version
- - Previous CUDA version
- - New CUDA version
- - Previous NCCL version
- - New NCCL version
-
-**Proposed fix**
-If you have a proposal for how to fix the issue state it here or link to a PR.
-
-**Additional context**
-Add any other context about the problem here.
diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml
deleted file mode 100644
index 58ba38e06..000000000
--- a/.github/workflows/stale.yml
+++ /dev/null
@@ -1,31 +0,0 @@
-# This workflow warns and then closes issues and PRs that have had no activity for a specified amount of time.
-#
-# You can adjust the behavior by modifying this file.
-# For more information, see:
-# https://github.com/actions/stale
-name: Mark stale issues and pull requests
-
-on:
-  schedule:
-  - cron: '15 18 * * *'
-
-jobs:
-  stale:
-
-    runs-on: ubuntu-latest
-    permissions:
-      issues: write
-      pull-requests: write
-
-    steps:
-    - uses: actions/stale@v5
-      with:
-        repo-token: ${{ secrets.GITHUB_TOKEN }}
-        days-before-stale: 60
-        stale-issue-message: 'Marking as stale. No activity in 60 days.'
-        stale-pr-message: 'Marking as stale. No activity in 60 days.'
-        stale-issue-label: 'stale'
-        stale-pr-label: 'stale'
-        remove-stale-when-updated: true
-        operations-per-run: 1000
-        days-before-close: -1
diff --git a/.gitignore b/.gitignore
deleted file mode 100644
index fd081062e..000000000
--- a/.gitignore
+++ /dev/null
@@ -1,14 +0,0 @@
-__pycache__
-*.so
-build
-.coverage_*
-*.egg-info
-*~
-slurm*
-logs
-.vscode
-local/
-.gitmodules
-wandb/
-onelogger.log
-onelogger.err
\ No newline at end of file
diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
deleted file mode 100644
index 3b2e4e150..000000000
--- a/.gitlab-ci.yml
+++ /dev/null
@@ -1,141 +0,0 @@
-workflow:
-  rules:
-    - if: $CI_PROJECT_NAMESPACE != "ADLR"
-      when: never
-    - if: $CI_COMMIT_BRANCH =~ /ci-/ && $CI_PIPELINE_SOURCE != "schedule"
-      when: never
-    - if: $CI_PIPELINE_SOURCE == "schedule"
-      auto_cancel:
-        on_new_commit: none
-    - if: $CI_PIPELINE_SOURCE == "web"
-    - if: $CI_COMMIT_REF_PROTECTED == "true"
-      variables:
-        FUNCTIONAL_TEST: 'no'
-    - if: $CI_MERGE_REQUEST_LABELS =~ /Run tests/ && $CI_MERGE_REQUEST_TARGET_BRANCH_SHA != ""
-      variables:
-        UNIT_TEST_REPEAT: 1
-        UNIT_TEST_TIMEOUT: 15
-        FUNCTIONAL_TEST: 'yes'
-        FUNCTIONAL_TEST_SCOPE: mr
-        FUNCTIONAL_TEST_REPEAT: 5
-        FUNCTIONAL_TEST_TIME_LIMIT: 2700
-        FUNCTIONAL_TEST_CLUSTER_A100: ''
-        FUNCTIONAL_TEST_CLUSTER_H100: ''
-        PUBLISH: 'no'
-    - if: $CI_MERGE_REQUEST_LABELS =~ /Run nightly/ && $CI_MERGE_REQUEST_TARGET_BRANCH_SHA != ""
-      variables:
-        UNIT_TEST_REPEAT: 1
-        UNIT_TEST_TIMEOUT: 15
-        FUNCTIONAL_TEST: 'yes'
-        FUNCTIONAL_TEST_SCOPE: nightly
-        FUNCTIONAL_TEST_REPEAT: 5
-        FUNCTIONAL_TEST_TIME_LIMIT: 2700
-        FUNCTIONAL_TEST_CLUSTER_A100: ''
-        FUNCTIONAL_TEST_CLUSTER_H100: ''
-        PUBLISH: 'no'
-    - if: $CI_MERGE_REQUEST_LABELS =~ /Run weekly/ && $CI_MERGE_REQUEST_TARGET_BRANCH_SHA != ""
-      variables:
-        UNIT_TEST_REPEAT: 1
-        UNIT_TEST_TIMEOUT: 15
-        FUNCTIONAL_TEST: 'yes'
-        FUNCTIONAL_TEST_SCOPE: weekly
-        FUNCTIONAL_TEST_REPEAT: 1
-        FUNCTIONAL_TEST_TIME_LIMIT: 9000
-        FUNCTIONAL_TEST_CLUSTER_A100: ''
-        FUNCTIONAL_TEST_CLUSTER_H100: ''
-        PUBLISH: 'no'
-    - if: $CI_PIPELINE_SOURCE == "merge_request_event" && $CI_MERGE_REQUEST_TARGET_BRANCH_SHA != ""
-      variables:
-        FUNCTIONAL_TEST: 'no'
-        PUBLISH: 'no'
-    - when: never
-  auto_cancel:
-    on_new_commit: interruptible
-    # on_job_failure: all
-
-stages:
-  - test
-  - functional_tests
-  - publish
-
-default:
-  interruptible: true
-  retry:
-    max: 2
-    when: runner_system_failure
-
-variables:
-  UNIT_TEST:
-    value: 'yes'
-    options:
-      - 'yes'
-      - 'no'
-    description: To run the funtional test suite
-  UNIT_TEST_REPEAT:
-    value: '1'
-    description: 'Number of repetitions'
-  UNIT_TEST_TIMEOUT:
-    value: '30'
-    description: Timeout (minutes) for Unit tests (all repeats)
-  FUNCTIONAL_TEST:
-    value: 'yes'
-    options:
-      - 'yes'
-      - 'no'
-    description: To run the funtional test suite
-  FUNCTIONAL_TEST_SCOPE:
-    value: 'mr'
-    options:
-      - 'mr'
-      - 'nightly'
-      - 'weekly'
-      - 'pre-release'
-      - 'release'
-    description: 'Testsuite to run (only for FUNCTIONAL_TEST=yes)'
-  FUNCTIONAL_TEST_REPEAT:
-    value: '5'
-    description: 'Number of repetitions per test'
-  FUNCTIONAL_TEST_TIME_LIMIT:
-    value: '2700'
-    description: 'Timeout in seconds per test'
-  FUNCTIONAL_TEST_CASES:
-    value: 'all'
-    description: "Comma-separated list of test_cases to run. Use 'all' to run the full suite."
-  FUNCTIONAL_TEST_CLUSTER_A100:
-    value: 'dgxa100_dracooci'
-    options:
-      - 'dgxa100_dracooci'
-      - 'dgxa100_dracooci-ord'
-    description: 'Cluster for A100 workloads'
-  FUNCTIONAL_TEST_CLUSTER_H100:
-    value: 'dgxh100_eos'
-    options:
-      - 'dgxh100_coreweave'
-      - 'dgxh100_eos'
-    description: 'Cluster for H100 workloads'
-  FUNCTIONAL_TEST_NAME:
-    description: 'Name of functional test run (only for pre-release and release)'
-  PUBLISH:
-    value: 'no'
-    options:
-      - 'yes'
-      - 'no'
-    description: Build and publish a wheel to PyPi
-  PUBLISH_SCOPE:
-    value: 'code-freeze'
-    options:
-      - 'code-freeze'
-      - 'release'
-    description: Type of publish (freeze or final release)
-
-  # CI wide variables
-  CI_MCORE_LTS_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/mcore_ci_lts
-  CI_MCORE_DEV_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/mcore_ci_dev
-  CI_NEMO_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/nemo_ci
-  UTILITY_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/mcore_utility
-
-include:
-  - .gitlab/stages/00.pre.yml
-  - .gitlab/stages/01.test.yml
-  - .gitlab/stages/02.functional-tests.yml
-  - .gitlab/stages/03.publish.yml
diff --git a/.gitlab/labeler-config.yml b/.gitlab/labeler-config.yml
deleted file mode 100644
index 3dc4001cd..000000000
--- a/.gitlab/labeler-config.yml
+++ /dev/null
@@ -1,33 +0,0 @@
-CI:
-- .gitlab-ci.yml
-- Dockerfile.ci.lts
-- Dockerfile.ci.dev
-- .github/**
-- .gitlab/**
-
-Datasets:
-- megatron/core/datasets/**
-
-BERT:
-- megatron/core/models/bert/**
-
-GPT:
-- megatron/core/models/gpt/**
-
-RETRO:
-- megatron/core/models/retro/**
-
-Dist-Ckpt:
-- megatron/core/dist_checkpointing
-
-Dist-Opt:
-- megatron/core/optimizer/distrib_optimizer 
-
-Inference:
-- megatron/core/inference
-
-MoE:
-- megatron/core/transformer/moe
-
-Tests:
-- tests/**
\ No newline at end of file
diff --git a/.gitlab/stages/00.pre.yml b/.gitlab/stages/00.pre.yml
deleted file mode 100644
index b5af2eeb8..000000000
--- a/.gitlab/stages/00.pre.yml
+++ /dev/null
@@ -1,199 +0,0 @@
-include:
-  - template: Security/Secret-Detection.gitlab-ci.yml
-
-.pre_rules:
-  rules:
-    - if: $CI_PIPELINE_SOURCE == 'merge_request_event' && $CI_MERGE_REQUEST_TARGET_BRANCH_PROTECTED != "true"
-      allow_failure: true
-      when: always
-    - if: $CI_PIPELINE_SOURCE == 'merge_request_event'
-    - when: never
-  stage: .pre
-
-.dind_rules:
-  image: docker:26.1.4-dind
-  variables:
-    DOCKER_HOST: unix:///var/run/docker.sock
-  before_script:
-    - docker system prune -a --filter "until=36h" -f || true
-    - echo "$NGC_API_KEY" | docker login nvcr.io -u '$oauthtoken' --password-stdin
-    - echo "$CI_REGISTRY_PASSWORD" | docker login $CI_REGISTRY -u $CI_REGISTRY_USER --password-stdin
-
-pre:mirror_to_github:
-  rules:
-    - if: '$CI_COMMIT_REF_PROTECTED == "true" && $CI_PIPELINE_SOURCE == "push"'
-    - when: never
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  stage: .pre
-  image: python:3.10
-  variables:
-    GIT_STRATEGY: 'clone'
-  script:
-    - git checkout $CI_COMMIT_BRANCH
-    - git remote add github https://ko3n1g:$GH_TOKEN@github.com/NVIDIA/Megatron-LM.git || true
-    - git push -u github $CI_COMMIT_BRANCH
-
-pre:create_ci_branches:
-  rules:
-    - if: '$CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH && $CI_PIPELINE_SOURCE == "push"'
-    - when: never
-  parallel:
-    matrix:
-      - branch: ci-unit-test-extended
-      - branch: ci-rebuild-mcore-nemo-image
-      - branch: ci-mr
-      - branch: ci-nightly
-      - branch: ci-weekly
-      - branch: ci-pre-release
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  stage: .pre
-  image: python:3.10
-  variables:
-    GIT_STRATEGY: 'clone'
-  script:
-    - git remote set-url origin "https://gitlab-ci-token:${PROJECT_ACCESS_TOKEN_MCORE}@${GITLAB_ENDPOINT}/adlr/megatron-lm.git"
-    - git switch --force-create $branch
-    - git push --force -u origin $branch
-
-pre:label_merge_request:
-  extends: [.pre_rules]
-  image: golang:1.22
-  tags:
-    - mcore-docker-node-small
-  before_script:
-    - git clone -b nv https://${GITLAB_ENDPOINT}/okoenig/gitlab-mr-labeler.git
-    - cd gitlab-mr-labeler
-    - go install .
-    - cd ..
-    - go install github.com/itchyny/gojq/cmd/gojq@latest
-    - |
-      echo LABELS=$(curl --header "PRIVATE-TOKEN: ${PROJECT_ACCESS_TOKEN_MCORE}" --url "https://${GITLAB_ENDPOINT}/api/v4/projects/${CI_PROJECT_ID}/merge_requests/${CI_MERGE_REQUEST_IID}" | gojq '.labels | join(",")') > labels
-  script:
-    - gitlab-mr-labeler -f .gitlab/labeler-config.yml -t ${PROJECT_ACCESS_TOKEN_MCORE} --debug true
-  after_script:
-    - |
-      source labels
-      curl --header "PRIVATE-TOKEN: ${PROJECT_ACCESS_TOKEN_MCORE}" --url "https://${GITLAB_ENDPOINT}/api/v4/projects/${CI_PROJECT_ID}/merge_requests/${CI_MERGE_REQUEST_IID}" --data-urlencode "add_labels=$LABELS" -X PUT
-
-pre:maybe_cherry_pick_commit:
-  rules:
-    - if: '$CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH && $CI_PIPELINE_SOURCE == "push"'
-    - when: never
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  stage: .pre
-  image: nentangso/alpine-git-curl-jq
-  variables:
-    GIT_STRATEGY: 'clone'
-  script:
-    - set -x
-    - set +e
-    - SHA=$(git rev-list --no-merges -n 1 HEAD)
-    - MESSAGE=$(git log -n 1 --pretty=format:%s $SHA)
-    - MR_ID=$(echo $MESSAGE | awk -F'!' '{print $2}' | awk '{print $1}' )
-    - git remote set-url origin "https://gitlab-ci-token:${PROJECT_ACCESS_TOKEN_MCORE}@${GITLAB_ENDPOINT}/$CI_PROJECT_NAMESPACE/megatron-lm.git"
-    - git config --global user.email "mcore-bot@nvidia.com"
-    - git config --global user.name "Mcore Bot"
-    - |
-      MR=$(curl --header "PRIVATE-TOKEN: ${PROJECT_ACCESS_TOKEN_MCORE}" --url "https://${GITLAB_ENDPOINT}/api/v4/projects/${CI_PROJECT_ID}/merge_requests/${MR_ID}")
-
-      LABELS=$(echo -E $MR | jq '.labels | join(",")' | tr -d '"')
-      AUTHOR_ID=$(echo -E $MR | jq '.author.id' | tr -d '"')
-      AUTHOR_NAME=$(echo -E $MR | jq '.author.username' | tr -d '"')
-      TITLE=$(echo -E $MR | jq '.title' | tr -d '"')
-      MILESTONE_ID=$(echo -E $MR | jq '.milestone.id' | tr -d '"')
-      TARGET_BRANCHES=$(echo "$LABELS" | grep -o 'core_[^,]*')
-
-      if [[ $TARGET_BRANCHES == "" ]]; then
-        echo Nothing to cherry pick
-        exit 0
-      fi
-
-      echo $TARGET_BRANCHES | while read -r RELEASE_BRANCH ; do
-        TARGET_BRANCH_EXISTS_OK=$([[ "$(git ls-remote --heads origin refs/heads/$RELEASE_BRANCH)" != "" ]] && echo true || echo false)
-
-        if [[ "$TARGET_BRANCH_EXISTS_OK" == "false" ]]; then
-          echo Release branch does not yet exist, will not  cherry-pick
-          continue
-        fi
-        
-        (
-          git fetch origin $RELEASE_BRANCH:$RELEASE_BRANCH
-          git switch --force-create cherry-pick-$MR_ID-$RELEASE_BRANCH $RELEASE_BRANCH
-          git cherry-pick $SHA
-          git push -u origin --force cherry-pick-$MR_ID-$RELEASE_BRANCH
-          git checkout ${CI_DEFAULT_BRANCH:-main}
-        )
-
-        CHERRYPICK_SUCCESSFUL=$?
-
-        if [[ $CHERRYPICK_SUCCESSFUL -eq 0 ]]; then
-          curl \
-            --header "PRIVATE-TOKEN: $PAT" \
-            --url https://${GITLAB_ENDPOINT}/api/v4/projects/${CI_PROJECT_ID}/merge_requests \
-            -d "source_branch=cherry-pick-$MR_ID-$RELEASE_BRANCH" \
-            -d "target_branch=$RELEASE_BRANCH" \
-            -d "title=Cherry pick \`$TITLE ($MR_ID)\` into \`$RELEASE_BRANCH\`" \
-            -d "labels=cherry-pick" \
-            -d "reviewer_ids=$AUTHOR_ID" \
-            -d "milestone_id=$MILESTONE_ID" \
-            -d "description=[🤖]: Hi @$AUTHOR_NAME 👋,<br><br>we've cherry picked \`$TITLE ($MR_ID)\` into \`$RELEASE_BRANCH\` for you! 🚀<br><br>Please review and approve this cherry pick by your convenience\!"
-
-        else
-          URL=https://${GITLAB_ENDPOINT}/ADLR/megatron-lm/-/merge_requests/$MR_ID
-
-          MESSAGE='{
-            "blocks": [
-              {
-                "type": "section",
-                "text": {
-                  "type": "mrkdwn",
-                  "text": "beep boop 🤖: Cherry-pick of <'$URL'|!'$MR_ID'> failed\ncc '$SLACK_ADMIN'"
-                }
-              }
-            ]
-          }'
-
-          curl -X POST -H "Content-type: application/json" --data "$MESSAGE" ${MCORE_NOTIFICATION_HOOK}
-
-        fi
-
-      done
-  interruptible: false
-
-pre:check_milestone:
-  extends: [.pre_rules]
-  image: badouralix/curl-jq
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  script:
-    - env
-    - |
-      MILESTONE=$(curl --header "PRIVATE-TOKEN: ${PROJECT_ACCESS_TOKEN_MCORE}" --url "https://${GITLAB_ENDPOINT}/api/v4/projects/${CI_PROJECT_ID}/merge_requests/${CI_MERGE_REQUEST_IID}" | jq '.milestone')
-    - |
-      if [[ "$MILESTONE" == "null" ]]; then
-        echo Please assign a Milestone to this MR!
-        exit 1
-      fi
diff --git a/.gitlab/stages/01.test.yml b/.gitlab/stages/01.test.yml
deleted file mode 100644
index 50d38fd7e..000000000
--- a/.gitlab/stages/01.test.yml
+++ /dev/null
@@ -1,611 +0,0 @@
-.test_rules:
-  rules:
-    - if: $UNIT_TEST == 'yes' && $CI_PIPELINE_SOURCE == 'merge_request_event' && $CI_MERGE_REQUEST_TARGET_BRANCH_PROTECTED != "true"
-      allow_failure: true
-      when: on_success
-    - when: on_success
-  stage: test
-
-include:
-  - template: Security/Secret-Detection.gitlab-ci.yml
-
-test:build_image:
-  extends: [.test_rules, .dind_rules]
-  tags:
-    - arch/amd64
-    - origin/jet-fleet
-    - env/prod
-    - ${TAG}
-  services:
-    - name: docker:24.0.5-dind
-      variables:
-        HEALTHCHECK_TCP_PORT: '2376'
-  timeout: 45m
-  parallel:
-    matrix:
-      - IMAGE: CI_MCORE_LTS_IMAGE
-        FILE: Dockerfile.ci.lts
-        BASE_IMAGE: nvcr.io/nvidia/pytorch:24.01-py3
-      - IMAGE: CI_MCORE_DEV_IMAGE
-        FILE: Dockerfile.ci.dev
-        BASE_IMAGE: nvcr.io/nvidia/pytorch:24.10-py3
-      - IMAGE: CI_NEMO_IMAGE
-        FILE: Dockerfile.ci.dev
-        BASE_IMAGE: nvcr.io/nvidian/nemo:nightly
-      - IMAGE: UTILITY_IMAGE
-        FILE: Dockerfile.linting
-        BASE_IMAGE: python:3.10
-  variables:
-    DOCKER_HOST: tcp://docker:2376
-    DOCKER_TLS_CERTDIR: '/certs'
-    DOCKER_TLS_VERIFY: 1
-    DOCKER_CERT_PATH: '$DOCKER_TLS_CERTDIR/client'
-    TAG: purpose/builder-large
-    STAGE: jet
-    MCORE_BACKWARDS_REF: core_r0.11.0
-  script:
-    - apk add bash
-    - |
-      bash -c '
-        set -x
-        env
-        eval "IMAGE=\$$IMAGE"
-        
-        docker context create tls-environment
-        docker buildx create --name container --driver=docker-container --use tls-environment
-
-        ADDITIONAL_PARAMS=()
-
-        if [[ "$CI_COMMIT_BRANCH" == "ci-rebuild-mcore-nemo-image" || "$CI_COMMIT_BRANCH" == "main" ]]; then
-          ADDITIONAL_PARAMS+=("--pull")
-          ADDITIONAL_PARAMS+=("--cache-to type=registry,ref=${IMAGE}-buildcache:main")
-        fi
-
-        if [[ "$CI_COMMIT_BRANCH" == "ci-nightly" ]]; then
-          ADDITIONAL_PARAMS+=("-t ${IMAGE}:nightly")
-        fi
-
-        echo $(git rev-parse HEAD)
-
-        DOCKER_BUILDKIT=1 docker build \
-          --secret id=JET_INDEX_URLS \
-          --secret id=LOGGER_INDEX_URL \
-          --target $STAGE \
-          -f $FILE \
-          -t ${IMAGE}:${CI_PIPELINE_ID} \
-          --builder=container \
-          --build-arg CACHEBUST=$(cat /proc/sys/kernel/random/uuid) \
-          --build-arg MCORE_REPO=${CI_REPOSITORY_URL} \
-          --build-arg MCORE_REF=$CI_COMMIT_SHA \
-          --build-arg MCORE_BACKWARDS_REF=$MCORE_BACKWARDS_REF \
-          --cache-to type=registry,ref=${IMAGE}-buildcache:${CI_PIPELINE_ID} \
-          --cache-to type=registry,ref=${IMAGE}-buildcache:${CI_MERGE_REQUEST_IID:-noop} \
-          --cache-from type=registry,ref=${IMAGE}-buildcache:main \
-          --cache-from type=registry,ref=${IMAGE}-buildcache:${CI_PIPELINE_ID} \
-          --cache-from type=registry,ref=${IMAGE}-buildcache:${CI_MERGE_REQUEST_IID:-noop} \
-          --build-arg FROM_IMAGE_NAME=$BASE_IMAGE \
-          --push \
-          ${ADDITIONAL_PARAMS[@]} .
-        '
-  retry:
-    max: 2
-
-test:unit_tests_configure:
-  extends: [.test_rules]
-  needs:
-    - test:build_image
-  image: ${UTILITY_IMAGE}:${CI_PIPELINE_ID}
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  before_script:
-    - git rm -r tests/test_utils/local_recipes || true
-    - git submodule add --force https://gitlab-ci-token:${CI_JOB_TOKEN}@${GITLAB_ENDPOINT}/ADLR/megatron-lm-convergence-tests.git tests/test_utils/local_recipes
-    - ls tests/test_utils/local_recipes
-  script:
-    - set -x
-    - |
-      A100_CLUSTER=$([[ "$FUNCTIONAL_TEST_CLUSTER_A100" != "" ]] && echo $FUNCTIONAL_TEST_CLUSTER_A100 || echo $DEFAULT_A100_CLUSTER)
-      H100_CLUSTER=$([[ "$FUNCTIONAL_TEST_CLUSTER_H100" != "" ]] && echo $FUNCTIONAL_TEST_CLUSTER_H100 || echo $DEFAULT_H100_CLUSTER)
-    - |
-      ARGS=(
-        "--scope unit-tests"
-        "--n-repeat ${UNIT_TEST_REPEAT}"
-        "--time-limit $(( UNIT_TEST_TIMEOUT * 60 ))"
-        "--test-cases all"
-        "--a100-cluster dgxa100_dracooci-ord"
-        "--h100-cluster dgxh100_coreweave"
-        "--h100-partition batch_short,batch"
-        "--container-image ${UTILITY_IMAGE}"
-        "--container-tag ${CI_PIPELINE_ID}"
-        "--dependent-job test:unit_tests_configure"
-        "--slurm-account ${CI_SLURM_ACCOUNT}"
-      )
-    - |
-      export PYTHONPATH=$(pwd)
-      python tests/test_utils/python_scripts/generate_jet_trigger_job.py \
-        ${ARGS[@]} \
-        --environment "lts" \
-        --tag "legacy" \
-        --output-path "unit-test-job-lts-legacy.yaml"
-    - |
-      export PYTHONPATH=$(pwd)
-      python tests/test_utils/python_scripts/generate_jet_trigger_job.py \
-        ${ARGS[@]} \
-        --environment "lts" \
-        --tag "latest" \
-        --output-path "unit-test-job-lts-latest.yaml"
-    - |
-      export PYTHONPATH=$(pwd)
-      python tests/test_utils/python_scripts/generate_jet_trigger_job.py \
-        ${ARGS[@]} \
-        --environment "dev" \
-        --tag "legacy" \
-        --output-path "unit-test-job-dev-legacy.yaml"
-    - |
-      export PYTHONPATH=$(pwd)
-      python tests/test_utils/python_scripts/generate_jet_trigger_job.py \
-        ${ARGS[@]} \
-        --environment "dev" \
-        --tag "latest" \
-        --output-path "unit-test-job-dev-latest.yaml"
-  rules:
-    - if: $UNIT_TEST == 'yes' && $CI_PIPELINE_SOURCE == 'merge_request_event' && $CI_MERGE_REQUEST_TARGET_BRANCH_PROTECTED != "true"
-      allow_failure: true
-      when: on_success
-    - if: $UNIT_TEST == 'yes' && $UNIT_TEST_REPEAT != '0'
-      when: on_success
-  artifacts:
-    paths:
-      - unit-test-job-dev-legacy.yaml
-      - unit-test-job-dev-latest.yaml
-      - unit-test-job-lts-legacy.yaml
-      - unit-test-job-lts-latest.yaml
-      - tests/test_utils/local_recipes
-
-.unit_tests_run:
-  needs:
-    - test:formatting
-    - test:copyright
-    - job: test:secret_detection
-      optional: true
-    - test:unit_tests_configure
-  extends: [.test_rules]
-  trigger:
-    include:
-      - artifact: unit-test-job-$ENVIRONMENT-$TAG.yaml
-        job: test:unit_tests_configure
-    strategy: depend
-  variables:
-    RO_API_TOKEN: $PAT
-    CONTAINER_TAG: $CI_PIPELINE_ID
-    CI_MCORE_LTS_IMAGE: $CI_MCORE_LTS_IMAGE
-    GITLAB_ENDPOINT: $GITLAB_ENDPOINT
-    PARENT_PIPELINE_ID: $CI_PIPELINE_ID
-  inherit:
-    variables: true
-  rules:
-    - if: $UNIT_TEST == 'yes' && $CI_PIPELINE_SOURCE == 'merge_request_event' && $CI_MERGE_REQUEST_TARGET_BRANCH_PROTECTED != "true"
-      allow_failure: true
-      when: on_success
-    - if: $UNIT_TEST == 'yes' && $UNIT_TEST_REPEAT != '0'
-      when: on_success
-
-test:unit_tests_pyt(DEV)_mcore(legacy):
-  extends: [.unit_tests_run]
-  variables:
-    ENVIRONMENT: dev
-    TAG: legacy
-
-test:unit_tests_pyt(LTS)_mcore(legacy):
-  extends: [.unit_tests_run]
-  variables:
-    ENVIRONMENT: lts
-    TAG: legacy
-
-test:unit_tests_pyt(DEV)_mcore(latest):
-  extends: [.unit_tests_run]
-  variables:
-    ENVIRONMENT: dev
-    TAG: latest
-
-test:unit_tests_pyt(LTS)_mcore(latest):
-  extends: [.unit_tests_run]
-  variables:
-    ENVIRONMENT: lts
-    TAG: latest
-
-test:notify_unit_tests:
-  extends: [.test_rules]
-  image: badouralix/curl-jq
-  needs:
-    - test:unit_tests_pyt(DEV)_mcore(latest)
-    - test:unit_tests_pyt(LTS)_mcore(latest)
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  script:
-    - apk add bash
-    - apk add --update coreutils
-    - env
-    - export WEBHOOK_URL=${MCORE_NOTIFICATION_HOOK}
-    - export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}
-    - export GITLAB_ENDPOINT
-    - export CONTEXT="unit-tests-extended"
-    - export DATE=$(date +"%Y-%m-%d")
-    - bash tests/test_utils/shell_scripts/notify.sh ${CI_PIPELINE_ID} "test:unit_tests_pyt"
-  artifacts:
-    when: always
-    paths:
-      - scripts
-  rules:
-    - if: $CI_PIPELINE_SOURCE == "schedule" && $CI_COMMIT_BRANCH == "ci-unit-test-extended"
-      when: always
-    - when: never
-
-test:docs_build:
-  extends: [.test_rules]
-  image: ${UTILITY_IMAGE}:${CI_PIPELINE_ID}
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  needs: [test:build_image]
-  script:
-    - cd ..
-    - rm -rf documentation && git clone https://gitlab-ci-token:${CI_JOB_TOKEN}@${GITLAB_ENDPOINT}/nemo-megatron-core-tme/documentation.git
-    - mv megatron-lm/ documentation/
-    - cd documentation/
-    - ./repo docs
-
-test:formatting:
-  extends: [.test_rules]
-  image: ${UTILITY_IMAGE}:${CI_PIPELINE_ID}
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  needs: [test:build_image]
-  variables:
-    GIT_STRATEGY: 'clone'
-  script:
-    - |
-      if [[ "$CI_PIPELINE_SOURCE" != "merge_request_event" ]]; then
-        exit 0
-      fi
-    - set +e
-    - git fetch origin main:main
-    - |
-      if [[ "$CI_MERGE_REQUEST_PROJECT_PATH" == "$CI_MERGE_REQUEST_SOURCE_PROJECT_PATH" ]]; then 
-        bash tools/autoformat.sh
-        set -e
-        git fetch origin $CI_MERGE_REQUEST_SOURCE_BRANCH_NAME
-        git checkout $CI_MERGE_REQUEST_SOURCE_BRANCH_NAME
-        git config --global user.email "mcore-bot@nvidia.com"
-        git config --global user.name "Mcore Bot"
-        git remote set-url origin "https://gitlab-ci-token:${PAT}@${GITLAB_ENDPOINT}/$CI_PROJECT_NAMESPACE/megatron-lm.git"
-        git add -A .
-        git commit -m "chore: Format files" || true
-        git push -u origin $CI_MERGE_REQUEST_SOURCE_BRANCH_NAME
-      fi
-    - env
-    - BASE_REF="$CI_MERGE_REQUEST_TARGET_BRANCH_NAME" CHECK_ONLY=true SKIP_DOCS=$([[ "$CI_MERGE_REQUEST_LABELS" == *"Skip docs"* ]] && echo "true" || echo "false") bash tools/autoformat.sh
-
-test:copyright:
-  extends: [.test_rules]
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  image: ${UTILITY_IMAGE}:${CI_PIPELINE_ID}
-  needs: [test:build_image]
-  script:
-    - git fetch origin main
-    - bash tools/copyright.sh
-
-# Override from template
-secret_detection:
-  rules:
-    - when: never
-
-# Inherit and modify template
-test:secret_detection:
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  extends: ['.secret-analyzer']
-  variables:
-    GIT_DEPTH: 0
-    SECRET_DETECTION_LOG_OPTIONS: ${CI_MERGE_REQUEST_DIFF_BASE_SHA}..${CI_COMMIT_SHA}
-  allow_failure: true
-  rules:
-    - if: $CI_PIPELINE_SOURCE == "merge_request_event"
-    - when: never
-  script:
-    - apk add jq
-    - /analyzer run
-    - |
-      if [[ $(cat gl-secret-detection-report.json | jq '.vulnerabilities | length > 0') == true ]]; then
-        echo "Atleast one vulnerability has been found"
-        cat gl-secret-detection-report.json | jq '.'
-        exit 1
-      fi
-
-test:pypi_build_wheel:
-  extends: [.test_rules]
-  image:
-    name: quay.io/pypa/manylinux_2_28_x86_64
-    entrypoint: ['']
-  services:
-    - name: docker:24.0.5-dind
-      variables:
-        HEALTHCHECK_TCP_PORT: '2376'
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/builder-small
-    - team/megatron
-  variables:
-    PUBLISH_DRYRUN: 'yes'
-    PY_ENV: pytorch_24.10
-  script:
-    - echo $PUBLISH_DRYRUN
-    - >
-      if [ "$PUBLISH_DRYRUN" = "yes" ]; then
-        PRE_RELEASE=$(sed -n "s/.*PRE_RELEASE = '\(.*\)'/\1/p" megatron/core/package_info.py)
-        sed -i "/^PRE_RELEASE/c\PRE_RELEASE = '${PRE_RELEASE}.dev$((RANDOM % 900000 + 100000))'" megatron/core/package_info.py 
-      fi
-
-
-    - /opt/python/cp310-cp310/bin/python -m build
-    - /opt/python/cp311-cp311/bin/python -m build
-    - auditwheel repair dist/*.whl
-    - rm -rf dist/*.whl
-
-    - pushd megatron/core
-    - EXPECTED_RELEASE_NUMBER=$(/opt/python/cp311-cp311/bin/python -c "import package_info; print(package_info.__version__)")
-    - popd
-    - echo "EXPECTED_RELEASE_NUMBER=$EXPECTED_RELEASE_NUMBER" | tee -a build.env
-  artifacts:
-    paths:
-      - megatron/core/package_info.py
-      - wheelhouse/
-      - dist/
-    reports:
-      dotenv: build.env
-
-test:pypi_test_wheel:
-  extends: [.test_rules]
-  image: 
-    name: python:3.11
-    entrypoint: ['']
-  needs: [test:pypi_build_wheel]
-  services:
-    - name: docker:24.0.5-dind
-      variables:
-        HEALTHCHECK_TCP_PORT: '2376'
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/builder-small
-    - team/megatron
-  variables:
-    PUBLISH_DRYRUN: 'yes'
-  script:
-    - rm -rf megatron
-    - pip install --no-cache-dir wheelhouse/*cp311*.whl
-
-    - RELEASE_NUMBER=$(python -c "from megatron import core; print(core.__version__)")
-    - >
-      echo "$EXPECTED_RELEASE_NUMBER" == "$RELEASE_NUMBER"
-
-
-    - test "$EXPECTED_RELEASE_NUMBER" == "$RELEASE_NUMBER"
-    - echo "RELEASE_NUMBER=$EXPECTED_RELEASE_NUMBER" | tee -a build.env
-  artifacts:
-    reports:
-      dotenv: build.env
-    paths:
-      - wheelhouse/
-      - dist/
-
-test:pypi_push_wheel:
-  extends: [.test_rules]
-  image: python:3.11
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  needs: [test:pypi_test_wheel]
-  variables:
-    PUBLISH_DRYRUN: 'yes'
-  timeout: 3m
-  script:
-    - >
-      if [ "$PUBLISH_DRYRUN" = "yes" ]; then
-        REPOSITORY=testpypi
-        export TWINE_USERNAME=$TWINE_TEST_USERNAME
-        export TWINE_PASSWORT=$TWINE_TEST_PASSWORD
-      else
-        REPOSITORY=pypi
-        export TWINE_USERNAME=$TWINE_PROD_USERNAME
-        export TWINE_PASSWORT=$TWINE_PROD_PASSWORD
-      fi
-
-    - ls -al dist/
-    - ls -al wheelhouse/
-    - pip install twine
-    - twine upload --verbose -u $TWINE_USERNAME -p $TWINE_PASSWORT --repository $REPOSITORY wheelhouse/* dist/*
-  
-  rules:
-    - if: $UNIT_TEST == 'yes' && $CI_PIPELINE_SOURCE == 'merge_request_event' && $CI_MERGE_REQUEST_TARGET_BRANCH_PROTECTED != "true"
-      allow_failure: true
-      when: on_success
-    - when: on_success
-      allow_failure: true
-
-test:gh_release:
-  extends: [.test_rules]
-  needs: [test:pypi_test_wheel]
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  image: badouralix/curl-jq
-  variables:
-    PUBLISH_DRYRUN: 'yes'
-  script:
-    - NAME="NVIDIA Megatron Core $RELEASE_NUMBER"
-    - IS_PRERELEASE=$([[ "$RELEASE_NUMBER" == *rc* ]] && echo "true" || echo "false")
-    - >
-      if [[ "$IS_PRERELEASE" == "true" ]]; then
-        DATE=$(date +"%Y-%m-%d")
-        CHANGELOG="Prerelease: $NAME ($DATE)"
-      else
-        CHANGELOG=$(awk '/^## '"$NAME"'/{flag=1; next} /^## /{flag=0} flag' CHANGELOG.md)
-        CHANGELOG=$(echo "$CHANGELOG" | sed '/./!d')
-      fi
-    - >
-      PAYLOAD=$(jq -nc \
-                  --arg TAG_NAME "v${RELEASE_NUMBER}" \
-                  --arg CI_COMMIT_SHA "$CI_COMMIT_SHA" \
-                  --arg NAME "$NAME" \
-                  --arg BODY "$CHANGELOG" \
-                  --argjson PRERELEASE "$IS_PRERELEASE" \
-                  '{
-                      "tag_name": $TAG_NAME,
-                      "target_commitish": $CI_COMMIT_SHA,
-                      "name": $NAME,
-                      "body": $BODY,
-                      "draft": false,
-                      "prerelease": $PRERELEASE,
-                      "generate_release_notes": false
-                  }'
-              )
-      echo -E "$PAYLOAD" > payload.txt
-    - cat payload.txt
-    - >
-      CMD=$(echo -E 'curl -L \
-        -X POST \
-        -H "Accept: application/vnd.github+json" \
-        -H "Authorization: Bearer '"$GH_TOKEN"'" \
-        -H "X-GitHub-Api-Version: 2022-11-28" \
-        https://api.github.com/repos/NVIDIA/Megatron-LM/releases \
-        -d @payload.txt
-      ')
-
-    - >
-      if [[ "$PUBLISH_DRYRUN" == "yes" ]]; then
-        echo -E "$CMD"
-      else
-        eval "$CMD"
-      fi
-
-
-test:notify_release:
-  needs: [test:pypi_test_wheel, test:pypi_push_wheel, test:gh_release]
-  extends: [.test_rules]
-  image: badouralix/curl-jq
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  variables:
-    PUBLISH_DRYRUN: 'yes'
-  script:
-    - URL="https://github.com/NVIDIA/Megatron-LM/releases/tag/core_r$RELEASE_NUMBER"
-    - >
-      MESSAGE='{
-          "blocks": [
-            {
-              "type": "section",
-              "text": {
-                "type": "mrkdwn",
-                    "text": "Releasebot 🤖: Megatron-Core released <'$URL'|core_r'"$RELEASE_NUMBER"'> 🚀"
-              }
-            }
-          ]
-        }'
-
-
-    - echo "$MESSAGE"
-    - >
-      CMD=$(echo curl \
-        -X POST \
-        -H "Content-type: application/json" \
-        --data "$MESSAGE" ${MCORE_NOTIFICATION_HOOK_MAIN}
-      )
-      
-      if [[ "$PUBLISH_DRYRUN" == "yes" ]]; then
-        echo "$CMD"
-      else
-        eval "$CMD"
-      fi
-
-test:generate_coverage_report:
-  extends: [.test_rules]
-  needs:
-    - test:unit_tests_pyt(DEV)_mcore(latest)
-    - test:unit_tests_pyt(LTS)_mcore(latest)
-  image: ${UTILITY_IMAGE}:${CI_PIPELINE_ID}
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  script:
-    - env
-    - export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}
-    - export GITLAB_ENDPOINT
-    - python tests/test_utils/python_scripts/download_coverage_results.py --pipeline-id ${CI_PIPELINE_ID}
-    - coverage combine --keep $(ls coverage_results/*/coverage_report)
-    - coverage report
-    - coverage xml
-  coverage: '/TOTAL.+ ([0-9]{1,3}%)/'
-  artifacts:
-    reports:
-      coverage_report:
-        coverage_format: cobertura
-        path: coverage.xml
-  rules:
-    - if: $UNIT_TEST == 'yes' && $CI_PIPELINE_SOURCE == 'merge_request_event' && $CI_MERGE_REQUEST_TARGET_BRANCH_PROTECTED != "true"
-      allow_failure: true
-      when: on_success
-    - if: $UNIT_TEST == 'yes' && $UNIT_TEST_REPEAT != '0'
-      when: on_success
\ No newline at end of file
diff --git a/.gitlab/stages/02.functional-tests.yml b/.gitlab/stages/02.functional-tests.yml
deleted file mode 100644
index ac13ee028..000000000
--- a/.gitlab/stages/02.functional-tests.yml
+++ /dev/null
@@ -1,188 +0,0 @@
-.functional_tests_rules:
-  stage: functional_tests
-  rules:
-    - if: $FUNCTIONAL_TEST == "yes" && ($CI_PIPELINE_SOURCE == 'merge_request_event' && $CI_MERGE_REQUEST_TARGET_BRANCH_PROTECTED != "true")
-      allow_failure: true
-    - if: $FUNCTIONAL_TEST == "yes"
-    - when: never
-
-default:
-  id_tokens:
-    VAULT_JWT_TOKEN:
-      aud: https://stg.vault.nvidia.com
-
-include:
-  - project: dl/jet/gitlab-templates
-    ref: main
-    file: downstreams.yml
-
-functional:configure:
-  needs:
-    - test:build_image
-    - job: test:unit_tests_pyt(DEV)_mcore(latest)
-      optional: true
-    - job: test:unit_tests_pyt(LTS)_mcore(latest)
-      optional: true
-  extends: [.functional_tests_rules]
-  image: ${UTILITY_IMAGE}:${CI_PIPELINE_ID}
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  before_script:
-    - git rm -r tests/test_utils/local_recipes || true
-    - git submodule add --force https://gitlab-ci-token:${CI_JOB_TOKEN}@${GITLAB_ENDPOINT}/ADLR/megatron-lm-convergence-tests.git tests/test_utils/local_recipes
-    - ls tests/test_utils/local_recipes
-  script:
-    - set -x
-    - |
-      A100_CLUSTER=$([[ "$FUNCTIONAL_TEST_CLUSTER_A100" != "" ]] && echo $FUNCTIONAL_TEST_CLUSTER_A100 || echo $DEFAULT_A100_CLUSTER)
-      H100_CLUSTER=$([[ "$FUNCTIONAL_TEST_CLUSTER_H100" != "" ]] && echo $FUNCTIONAL_TEST_CLUSTER_H100 || echo $DEFAULT_H100_CLUSTER)
-    - |
-      RECORD_CHECKPOINTS=$([[ "$CI_MERGE_REQUEST_LABELS" == *"Record checkpoints"* ]] && echo "true" || echo "false")
-    - |
-      if [[ "$FUNCTIONAL_TEST_SCOPE" == "release" || "$FUNCTIONAL_TEST_SCOPE" == "pre-release" ]]; then
-        FUNCTIONAL_TEST_NAME=$(eval echo $FUNCTIONAL_TEST_NAME)
-        RELEASE_ARGS=(
-          "--run-name"
-          $FUNCTIONAL_TEST_NAME
-          "--wandb-experiment"
-          $(echo $FUNCTIONAL_TEST_NAME | tr '/' '-')
-        )
-      else
-        RELEASE_ARGS=()
-      fi
-    - |
-      ARGS=(
-        "--scope $FUNCTIONAL_TEST_SCOPE"
-        "--n-repeat $FUNCTIONAL_TEST_REPEAT"
-        "--time-limit $FUNCTIONAL_TEST_TIME_LIMIT"
-        "--test-cases $FUNCTIONAL_TEST_CASES"
-        "--a100-cluster $A100_CLUSTER"
-        "--h100-cluster $H100_CLUSTER"
-        "--container-image ${UTILITY_IMAGE}"
-        "--container-tag ${CI_PIPELINE_ID}"
-        "--dependent-job functional:configure"
-        "--record-checkpoints ${RECORD_CHECKPOINTS}"
-        "--slurm-account ${CI_SLURM_ACCOUNT}"
-      )
-    - |
-      export PYTHONPATH=$(pwd)
-      python tests/test_utils/python_scripts/generate_jet_trigger_job.py \
-        ${ARGS[@]} \
-        --environment dev \
-        --output-path "functional-test-job-dev.yaml" \
-        ${RELEASE_ARGS[@]}
-    - |
-      export PYTHONPATH=$(pwd)
-      python tests/test_utils/python_scripts/generate_jet_trigger_job.py \
-        ${ARGS[@]} \
-        --environment lts \
-        --output-path "functional-test-job-lts.yaml" \
-        ${RELEASE_ARGS[@]}
-  artifacts:
-    paths:
-      - functional-test-job-lts.yaml
-      - functional-test-job-dev.yaml
-      - tests/test_utils/local_recipes
-
-.run:
-  stage: functional_tests
-  needs: [functional:configure]
-  extends: [.functional_tests_rules]
-  trigger:
-    include:
-      - artifact: functional-test-job-$ENVIRONMENT.yaml
-        job: functional:configure
-    strategy: depend
-  variables:
-    RO_API_TOKEN: $PAT
-    CONTAINER_TAG: $CI_PIPELINE_ID
-    CI_MCORE_LTS_IMAGE: $CI_MCORE_LTS_IMAGE
-    GITLAB_ENDPOINT: $GITLAB_ENDPOINT
-    PARENT_PIPELINE_ID: $CI_PIPELINE_ID
-  inherit:
-    variables: true
-
-functional:run_lts:
-  extends: [.run]
-  variables:
-    ENVIRONMENT: lts
-
-functional:run_dev:
-  extends: [.run]
-  variables:
-    ENVIRONMENT: dev
-
-functional:run_nemo:
-  extends: [.functional_tests_rules]
-  trigger:
-    project: 'dl/joc/nemo-ci'
-    branch: main-mirror
-    strategy: depend
-  inherit:
-    variables: true
-  variables:
-    MCORE_COMMIT: $CI_COMMIT_SHA
-    TEST_LLM_MODULE: 'True'
-    TEST_ALIGNER_MODULE: 'False'
-    TEST_DATA_CURATOR_MODULE: 'False'
-    TESTS_TO_RUN_ON_THIS_COMMIT: nightly
-  rules:
-    - if: $FUNCTIONAL_TEST == "yes"
-      when: manual
-      allow_failure: true
-    - when: never
-
-functional:notify:
-  extends: [.functional_tests_rules]
-  image: badouralix/curl-jq
-  needs:
-    - functional:run_lts
-    - functional:run_dev
-  tags:
-    - mcore-docker-node-small
-  variables:
-    WEBHOOK_URL: ${MCORE_NOTIFICATION_HOOK}
-    RO_API_TOKEN: ${PROJECT_ACCESS_TOKEN_MCORE}
-    CONTEXT: $FUNCTIONAL_TEST_SCOPE
-  script:
-    - apk add bash
-    - apk add --update coreutils
-    - env
-    - export WEBHOOK_URL=${MCORE_NOTIFICATION_HOOK}
-    - export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}
-    - export GITLAB_ENDPOINT
-    - export CONTEXT=$FUNCTIONAL_TEST_SCOPE
-    - export DATE=$(date +"%Y-%m-%d")
-    - bash tests/test_utils/shell_scripts/notify.sh ${CI_PIPELINE_ID} "functional:run_"
-  artifacts:
-    when: always
-    paths:
-      - scripts
-  rules:
-    - if: $CI_PIPELINE_SOURCE == "schedule" && $FUNCTIONAL_TEST == "yes"
-      when: always
-    - when: never
-
-functional:download_golden_values:
-  extends: [.functional_tests_rules]
-  image: ${UTILITY_IMAGE}:${CI_PIPELINE_ID}
-  tags:
-    - mcore-docker-node-small
-  script:
-    - env
-    - export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}
-    - export GITLAB_ENDPOINT
-    - python tests/test_utils/python_scripts/download_golden_values.py --pipeline-id ${CI_PIPELINE_ID}
-  artifacts:
-    paths:
-      - tests/
-  rules:
-    - if: $FUNCTIONAL_TEST == "yes"
-      when: manual
-      allow_failure: true
-    - when: never
diff --git a/.gitlab/stages/03.publish.yml b/.gitlab/stages/03.publish.yml
deleted file mode 100644
index 48ea9bfbf..000000000
--- a/.gitlab/stages/03.publish.yml
+++ /dev/null
@@ -1,126 +0,0 @@
-.publish_common_freeze:
-  stage: publish
-  rules:
-    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH && $PUBLISH == "yes" && $PUBLISH_SCOPE == "code-freeze"
-      when: manual
-    - when: never
-
-.publish_common_release:
-  stage: publish
-  rules:
-    - if: $CI_COMMIT_BRANCH =~ /^core_r/ && $PUBLISH == "yes" && $PUBLISH_SCOPE == "release"
-      when: manual
-    - if: $PUBLISH == "yes" && $PUBLISH_SCOPE == "release"
-      when: manual
-      variables:
-        PUBLISH_DRYRUN: 'yes'
-    - when: never
-
-publish:release_branch:
-  extends: [.publish_common_freeze]
-  image: ${CI_MCORE_LTS_IMAGE}:${CI_PIPELINE_ID}
-  needs: [test:build_image]
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  variables:
-    GIT_STRATEGY: 'none'
-  script:
-    - git fetch origin $CI_DEFAULT_BRANCH
-    - git config --global user.email "mcore-bot@nvidia.com"
-    - git config --global user.name "Mcore Bot"
-    - git remote set-url origin "https://gitlab-ci-token:${PAT}@${GITLAB_ENDPOINT}/$CI_PROJECT_NAMESPACE/megatron-lm.git"
-    - sed -i "/^PRE_RELEASE/c\PRE_RELEASE = ''" megatron/core/package_info.py
-    - VERSION=$(python -c "from megatron import core; print(core.__version__)")
-    - RELEASE_BRANCH=core_r$VERSION
-    - git switch --force-create $RELEASE_BRANCH origin/$CI_DEFAULT_BRANCH
-    - |
-      MESSAGE='{
-        "blocks": [
-          {
-            "type": "section",
-            "text": {
-              "type": "mrkdwn",
-              "text": "Releasebot 🤖: Megatron Core has been frozen 🎉 to branch `'"$RELEASE_BRANCH"'`"
-            }
-          }
-        ]
-      }'
-    - >
-      curl -X POST -H "Content-type: application/json" --data "$MESSAGE" ${MCORE_NOTIFICATION_HOOK_MAIN}
-
-
-    - git switch --force-create bot/chore/bump-version
-    - git add megatron/core/package_info.py
-    - >
-      git commit -m "chore: adjust version version"
-
-
-    - git push -u origin bot/chore/bump-version
-    - >
-      curl \
-        --header "PRIVATE-TOKEN: $PAT" \
-        --url https://${GITLAB_ENDPOINT}/api/v4/projects/${CI_PROJECT_ID}/merge_requests \
-        -d "source_branch=bot/chore/bump-version" \
-        -d "target_branch=$RELEASE_BRANCH" \
-        -d "title=chore: Fix version of \`$RELEASE_BRANCH\`" \
-        -d "description=[🤖]: Hi @okoenig 👋,<br><br>we've adjusted the version number of \`$RELEASE_BRANCH\` for you! 🚀<br><br>Please review and approve this cherry pick by your convenience\!"
-
-publish:pypi_build_wheel:
-  extends: [test:pypi_build_wheel, .publish_common_release]
-  dependencies: []
-  variables:
-    PUBLISH_DRYRUN: 'no'
-
-publish:pypi_test_wheel:
-  extends: [test:pypi_test_wheel, .publish_common_release]
-  needs: [publish:pypi_build_wheel]
-  variables:
-    PUBLISH_DRYRUN: 'no'
-
-publish:pypi_push_wheel:
-  extends: [test:pypi_push_wheel, .publish_common_release]
-  needs: [publish:pypi_test_wheel]
-  dependencies: [publish:pypi_test_wheel]
-  variables:
-    PUBLISH_DRYRUN: 'no'
-
-publish:gh_release:
-  extends: [test:gh_release, .publish_common_release]
-  dependencies: [publish:pypi_test_wheel]
-  needs: [publish:pypi_test_wheel]
-  variables:
-    PUBLISH_DRYRUN: 'no'
-
-publish:notify_release:
-  needs: [publish:pypi_push_wheel, publish:gh_release]
-  extends: [test:notify_release, .publish_common_release]
-  variables:
-    PUBLISH_DRYRUN: 'no'
-
-publish:docs:
-  extends: [.publish_common_release]
-  image: ${UTILITY_IMAGE}:${CI_PIPELINE_ID}
-  tags: 
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  script:
-    - cd ..
-    - rm -rf documentation && git clone https://gitlab-ci-token:${PROJECT_ACCESS_TOKEN_MCORE}@${GITLAB_ENDPOINT}/nemo-megatron-core-tme/documentation.git
-    - cd documentation/megatron-lm
-      git fetch origin '+refs/merge-requests/*:refs/remotes/merge-requests/*'
-    - git fetch origin $CI_COMMIT_SHA
-    - git checkout $CI_COMMIT_SHA
-    - cd ..
-    - git add megatron-lm
-    - >
-      git commit -m 'feat: Bump mcore'
-    - git push
diff --git a/.pylintrc b/.pylintrc
deleted file mode 100644
index 865f48384..000000000
--- a/.pylintrc
+++ /dev/null
@@ -1,13 +0,0 @@
-[MAIN]
-ignore-paths=tests
-max-line-length=100
-
-[MESSAGES CONTROL]
-disable=all
-
-enable=C0115,C0116,W0611,C0301,E0606
-# C0115: missing-class-docstring
-# C0116: missing-function-docstring
-# W0611: unused-import
-# C0301: line-too-long
-# E0606: possibly-used-before-assignment
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/cuda_patch/__init__.py b/megatron-lm-musa-patch/cuda_patch/__init__.py
new file mode 100644
index 000000000..d9bd50acd
--- /dev/null
+++ b/megatron-lm-musa-patch/cuda_patch/__init__.py
@@ -0,0 +1,28 @@
+import sys
+import torch
+import torch.utils
+import torch.utils.data
+
+from . import transformer_config
+from . import training
+from . import moe_utils
+from . import multi_latent_attention
+from . import router
+from . import arguments
+from . import theoretical_memory_usage
+
+from . import fused_layer_norm
+
+
+from . import training
+def py_patch():
+    if sys.version_info >= (3.9, 0):
+        return
+    import math
+    def lcm(a, b):
+        return abs(a * b) // math.gcd(a, b)
+    math.lcm = lcm
+    return
+
+# Apply patch
+py_patch()
diff --git a/megatron-lm-musa-patch/cuda_patch/arguments.py b/megatron-lm-musa-patch/cuda_patch/arguments.py
new file mode 100644
index 000000000..5b1683571
--- /dev/null
+++ b/megatron-lm-musa-patch/cuda_patch/arguments.py
@@ -0,0 +1,140 @@
+import dataclasses
+import torch
+import torch.nn.functional as F
+
+import megatron.training.arguments
+from megatron.training.activations import squared_relu
+from .transformer_config import TransformerConfig, MLATransformerConfig
+moe_freq_type = megatron.training.arguments.moe_freq_type
+
+def _add_moe_args(parser):
+    group = parser.add_argument_group(title="moe")
+    # General arguments
+    group.add_argument('--expert-model-parallel-size', type=int, default=1,
+                       help='Degree of expert model parallelism.')
+    group.add_argument('--expert-tensor-parallel-size', type=int, default=None,
+                       help='Degree of expert model parallelism. Default is None, which will be set to the value of --tensor-model-paralle-size.')
+    group.add_argument('--num-experts', type=int, default=None,
+                       help='Number of Experts in MoE (None means no MoE)')
+    group.add_argument('--moe-layer-freq', type=moe_freq_type, default=1,
+                       help='Frequency between MoE layers and Dense layers. Accepts either: '
+                            '- An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers '
+                            '- A string containing a Python list expression that defines a custom pattern, e.g.: '
+                            '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] '
+                            'where 1 indicates an expert layer and 0 indicates a dense layer. '
+                            'Examples: "([0]+[1]*23)": 1 dense layer followed by 23 experts layers, '
+                            '"([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice.')
+    group.add_argument('--moe-ffn-hidden-size', type=int, default=None,
+                       help='The hidden size of each expert\'s feed-forward network (ffn). '
+                       'If not specified, defaults to the ffn_hidden_size.')
+    group.add_argument('--moe-shared-expert-intermediate-size', type=int, default=None,
+                       help='Shared expert total ffn hidden size. '
+                       'It should be equal to "num_shared_experts * ffn_size_of_each_shared_expert" if there are multiple shared experts. '
+                       'None means no shared expert.')
+    group.add_argument('--moe-shared-expert-overlap', action='store_true',
+                       help='Enable overlapping between shared expert computations and dispatcher communications. '
+                       'Without this, the shared epxerts execute after the routed experts. '
+                       'Only effective when moe-shared-expert-intermediate-size is set.')
+    group.add_argument('--moe-grouped-gemm', action='store_true',
+                       help='When there are multiple experts per rank, launch multiple local GEMM kernels in multiple streams to improve the utilization and performance with GroupedLinear in TransformerEngine.')
+    # Router arguments
+    group.add_argument('--moe-router-load-balancing-type', type=str,
+                       choices=['aux_loss', 'seq_aux_loss', 'sinkhorn', 'none'],
+                       default='aux_loss',
+                       help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".')
+    group.add_argument('--moe-router-topk', type=int, default=2,
+                       help='Number of experts to route to for each token. The default is 2.')
+    group.add_argument('--moe-router-pre-softmax', action='store_true',
+                       help='Enable pre-softmax routing for MoE, which means softmax is before the top-k selection. By default, softmax is done after top-k.')
+    group.add_argument('--moe-router-topk-limited-devices', type=int, default=None,
+                       help='Number of expert parallel ranks to consider for each token during routing. Perform top-k routing on a subset of expert parallel ranks by first selecting N ranks for each token, then conducting top-k selection among experts on these devices. Default is None, which means no limited devices.')
+    group.add_argument('--moe-router-topk-scaling-factor', type=float, default=None,
+                       help='Scaling factor for routing score in top-k selection, only works when --moe-router-pre-softmax enabled. Defaults to None, which means no scaling.')
+    group.add_argument('--moe-use-legacy-grouped-gemm', action='store_true',
+                       help='Use legacy GroupedMLP rather than TEGroupedMLP. Note: The legacy one will be deprecated soon.')
+    group.add_argument('--moe-aux-loss-coeff', type=float, default=0.0,
+                       help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.')
+    group.add_argument('--moe-device-level-aux-loss-coeff', type=float, default=0.0,
+                       help='Scaling coefficient for the device-level aux loss')
+    group.add_argument('--moe-comm-aux-loss-coeff', type=float, default=0.0,
+                       help='Scaling coefficient for the communication aux loss')
+    group.add_argument('--moe-z-loss-coeff', type=float, default=None,
+                       help='Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended.')
+    group.add_argument('--moe-input-jitter-eps', type=float, default=None,
+                       help='Add noise to the input tensor by applying jitter with a specified epsilon value.')
+    group.add_argument('--moe-token-dispatcher-type', type=str,
+                       choices=['allgather', 'alltoall', 'alltoall_seq'],
+                       default='allgather',
+                       help="The type of token dispatcher to use. The default is 'allgather'. Options are 'allgather', 'alltoall' and 'alltoall_seq'. We recommend using 'alltoall' when applying expert parallelism. For more information, please refer to the documentation in core/moe/README.")
+    group.add_argument('--moe-per-layer-logging', action='store_true',
+                       help='Enable per-layer logging for MoE, currently supports auxiliary loss and z loss.')
+    # Token dropping arguments
+    group.add_argument('--moe-expert-capacity-factor', type=float, default=None,
+                       help='The capacity factor for each expert, None means no token will be dropped.')
+    group.add_argument('--moe-device-level-capacity', action='store_true',
+                       help='Whether to consider the expert capacity of a group together')
+    group.add_argument('--moe-pad-expert-input-to-capacity', action='store_true',
+                       help='Pads the input for each expert to match the expert capacity length, effective only after the --moe-expert-capacity-factor is set.')
+    group.add_argument('--moe-token-drop-policy', type=str, default='probs', choices=['probs', 'position'],
+                       help='The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with the lowest probabilities will be dropped. If "position", tokens at the end of each batch will be dropped.')
+    group.add_argument('--moe-layer-recompute', action='store_true',
+                       help='Enable checkpointing for moe_layer, should be used when memory is not sufficient.')
+    group.add_argument('--moe-extended-tp', action='store_true',
+                       help='Deprecated. Use --expert-tensor-parallel-size instead.')
+    group.add_argument('--moe-use-upcycling', action='store_true',
+                       help='Load a checkpoint of a dense model, convert it into an MoE model, and save the converted model to the path specified by --save. '
+                       'Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.')
+
+    return parser
+
+
+def core_transformer_config_from_args(args, config_class=None):
+
+    # Config class.
+    config_class = config_class or TransformerConfig
+
+    if args.multi_latent_attention:
+        config_class = MLATransformerConfig
+
+    # Translate args to core transformer configuration
+    kw_args = {}
+    for f in dataclasses.fields(config_class):
+        if hasattr(args, f.name):
+            kw_args[f.name] = getattr(args, f.name)
+    kw_args['persist_layer_norm'] = not args.no_persist_layer_norm
+    kw_args['layernorm_zero_centered_gamma'] = args.apply_layernorm_1p
+    kw_args['layernorm_epsilon'] = args.norm_epsilon
+    kw_args['deallocate_pipeline_outputs'] = True
+    kw_args['pipeline_dtype'] = args.params_dtype
+    kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm
+    kw_args['num_moe_experts'] = args.num_experts
+    kw_args['rotary_interleaved'] = args.rotary_interleaved
+    kw_args['first_pipeline_num_layers']= args.decoder_first_pipeline_num_layers
+    kw_args['last_pipeline_num_layers']= args.decoder_last_pipeline_num_layers
+    if args.swiglu:
+        kw_args['activation_func'] = F.silu
+        kw_args['gated_linear_unit'] = True
+        kw_args['bias_activation_fusion'] = args.bias_swiglu_fusion
+    else:
+        kw_args['bias_activation_fusion'] = args.bias_gelu_fusion
+    if args.squared_relu:
+        assert not args.swiglu
+        kw_args['activation_func'] = squared_relu
+    if args.init_method_xavier_uniform:
+        kw_args['init_method'] = torch.nn.init.xavier_uniform_
+        kw_args['scaled_init_method'] = torch.nn.init.xavier_uniform_
+    if args.group_query_attention:
+        kw_args['num_query_groups'] = args.num_query_groups
+    else:
+        kw_args['num_query_groups'] = None
+    kw_args['config_logger_dir'] = args.config_logger_dir
+
+    if len(args.cp_comm_type) == 1:
+        kw_args['cp_comm_type'] = args.cp_comm_type[0]
+
+    # Return config.
+    return config_class(**kw_args)
+
+
+megatron.training.arguments._add_moe_args = _add_moe_args
+megatron.training.arguments.core_transformer_config_from_args = core_transformer_config_from_args
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/cuda_patch/fused_layer_norm.py b/megatron-lm-musa-patch/cuda_patch/fused_layer_norm.py
new file mode 100644
index 000000000..bfc25b3c1
--- /dev/null
+++ b/megatron-lm-musa-patch/cuda_patch/fused_layer_norm.py
@@ -0,0 +1,106 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+import numbers
+
+import torch
+from torch import Tensor
+from torch.nn import init
+from torch.nn.parameter import Parameter
+
+from megatron.core.transformer import TransformerConfig
+
+
+
+class FusedLayerNorm(torch.nn.Module):
+
+    """Layer Norm, fused into a single CUDA kernel.
+
+    Args:
+      hidden_size (int): Transformer hidden dimension.
+
+      eps (float): Epsilon added to denominator, for numerical stability.
+
+      persist_layer_norm (bool): Use persistent fused layer norm kernel.
+      This kernel supports only a set of hidden sizes. Please
+      check persist_ln_hidden_sizes if your hidden size is supported.
+
+      zero_centered_gamma (bool): Adjust LayerNorm weights such that they are
+      centered around zero. This improves numerical stability.
+
+      config (TransformerConfig): Transformer config. Include to match custom
+      layer norm interfaces.
+
+      normalization (str): Normalization type, used for Transformer Engine.
+      Must equal 'LayerNorm' here.
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        hidden_size: int,
+        eps: float = 1e-5,
+        persist_layer_norm: bool = True,
+        zero_centered_gamma: bool = False,
+        normalization: str = "LayerNorm",  # included to match TE interface
+    ):
+        super().__init__()
+        print("use FusedLayerNorm")
+
+        self.config = config
+
+        self.zero_centered_gamma = self.config.layernorm_zero_centered_gamma
+
+        if self.config.normalization == "LayerNorm":
+            self.norm_impl = torch.layer_norm
+        elif self.config.normalization == "RMSNorm":
+            def naive_rms_norm(hidden_states, hidden_size, weight, eps):
+                variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+                hidden_states = hidden_states * torch.rsqrt(variance + eps)
+                # convert into half-precision if necessary
+                if self.weight.dtype in [torch.float16, torch.bfloat16]:
+                    hidden_states = hidden_states.to(self.weight.dtype)
+                hidden_states = weight * hidden_states
+                return hidden_states
+            self.norm_impl = naive_rms_norm
+        else:
+            raise ValueError(f'({self.config.normalization}) is not supported in FusedLayerNorm')
+
+        if isinstance(hidden_size, numbers.Integral):
+            hidden_size = (hidden_size,)
+        # self.hidden_size = torch.Size(hidden_size)
+        self.hidden_size = hidden_size
+        self.eps = eps
+        self.weight = Parameter(torch.Tensor(*hidden_size))
+        self.bias = Parameter(torch.Tensor(*hidden_size)) if self.config.normalization == "LayerNorm" else None
+        self.reset_parameters()
+        self.sequence_parallel = self.config.sequence_parallel
+
+
+        # set sequence parallelism flag on weight and bias parameters
+        setattr(self.weight, 'sequence_parallel', self.sequence_parallel)
+        if self.config.normalization == "LayerNorm":
+            setattr(self.bias, 'sequence_parallel', self.sequence_parallel)
+
+    def reset_parameters(self):
+
+        if self.zero_centered_gamma:
+            init.zeros_(self.weight)
+            if self.config.normalization == "LayerNorm":
+                init.zeros_(self.bias)
+        else:
+            init.ones_(self.weight)
+            if self.config.normalization == "LayerNorm":
+                init.zeros_(self.bias)
+
+    def forward(self, input: Tensor) -> Tensor:
+
+        weight = self.weight + 1 if self.zero_centered_gamma else self.weight
+        if self.config.normalization == "LayerNorm":
+            output = self.norm_impl(input, self.hidden_size, weight, self.bias, self.eps)
+        else:
+            output = self.norm_impl(input, self.hidden_size, weight, self.eps)
+
+        return output
+
+import megatron.core.fusions.fused_layer_norm
+megatron.core.fusions.fused_layer_norm.FusedLayerNorm = FusedLayerNorm
diff --git a/megatron-lm-musa-patch/cuda_patch/moe_utils.py b/megatron-lm-musa-patch/cuda_patch/moe_utils.py
new file mode 100644
index 000000000..ef250a60d
--- /dev/null
+++ b/megatron-lm-musa-patch/cuda_patch/moe_utils.py
@@ -0,0 +1,299 @@
+
+import math
+from typing import Optional
+
+import torch
+
+from megatron.core import parallel_state
+import megatron.core.transformer.moe.moe_utils
+get_capacity = megatron.core.transformer.moe.moe_utils.get_capacity
+device_limited_topk = megatron.core.transformer.moe.moe_utils.device_limited_topk
+
+
+def node_limited_topk(
+    scores: torch.Tensor,
+    topk: int,
+    num_tokens: int,
+    num_experts: int,
+    moe_router_topk_limited_devices: int,
+    num_node_group: int=None,
+):
+    """Perform top-k routing on a subset of expert parallel ranks.
+
+    Selects N ranks for each token, then conducts top-k selection among experts on these node.
+    See DeepSeek-V3 technical report for details.
+
+    Args:
+        scores (torch.Tensor): Softmax scores from the router.
+        topk (int): The number of experts to select for each token.
+        num_tokens (int): The number of tokens.
+        num_experts (int): The number of experts.
+        moe_router_topk_limited_devices (int): Number of expert parallel ranks to consider for
+            each token during routing. None means no device limitation.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Probs and indices tensor.
+    """
+
+    # Organize the experts into groups
+    if num_node_group is None:
+        ep_size = (
+            parallel_state.get_expert_model_parallel_world_size()
+        )  # num_node_group equals to expert parallel size/8
+        assert ep_size % 8 == 0, f"ep_size should be multiple of 8, but get {ep_size}"
+        num_node_group = ep_size // 8
+    node_k = topk // moe_router_topk_limited_devices #each token select node according to the sum of the highest K/M affinity scores
+    group_scores = (
+                scores.view(num_tokens, num_node_group, -1).topk(node_k, dim=-1)[0].sum(dim = -1)
+            )  # [n, n_group]
+    group_idx = torch.topk(
+                group_scores, k=moe_router_topk_limited_devices, dim=-1, sorted=False
+            )[
+                1
+            ]  # [n, moe_router_topk_limited_devices]
+    group_mask = torch.zeros_like(group_scores)  # [n, n_group]
+    group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(num_tokens, num_node_group, num_experts // num_node_group)
+        .reshape(num_tokens, -1)
+    )  # [n, e]
+    masked_scores = scores.masked_fill(~score_mask.bool(), 0.0)  # [n, e]
+    _, top_indices = torch.topk(masked_scores, k=topk, dim=-1)
+    return top_indices
+
+
+def sequence_load_balancing_loss_func(
+    probs: torch.Tensor,
+    routing_map: torch.Tensor,
+    tokens_per_expert: torch.Tensor,
+    batch_size: int,
+    seq_length: int,
+    topk: int,
+    moe_aux_loss_coeff: float,
+    moe_device_level_aux_loss_coeff: float=None,
+    moe_comm_aux_loss_coeff: float=None,
+    moe_router_topk_limited_devices: float=None,
+    moe_complementary_seq_aux_loss: bool=False,
+    sequence_partition_group=None,
+):
+    """
+    Calculate the auxiliary loss in sequence-level by computing the loss for each individual sample.
+    Refer to the DeepSeek-V2 huggingface repo
+    (https://huggingface.co/deepseek-ai/DeepSeek-V2) for details.
+    """
+    num_sub_sequence = 1
+
+    # If the sequence is partitioned by certain parallelism strategies like Sequence Parallelism
+    # or Context Parallelism, compute the gradient of the auxiliary loss with respect to the full
+    # sequence.
+    if sequence_partition_group is not None:
+        # We can keep `aggregated_probs_per_expert` local since we don't need the gradient for
+        # `tokens_per_expert`, saving one allreduce operation for `aggregated_probs_per_expert`.
+        num_sub_sequence = torch.distributed.get_world_size(sequence_partition_group)
+        torch.distributed.all_reduce(tokens_per_expert, group=sequence_partition_group)
+
+    assert num_sub_sequence == 1, "Do not support sequence aux loss in sequence partition case"
+
+    num_experts = probs.shape[1]
+
+    probs_for_aux_loss = probs.view(seq_length, batch_size, -1)
+    cost_coeff = routing_map.view(seq_length, batch_size, -1).sum(dim=0).float()
+    cost_coeff.div_(seq_length * topk / num_experts)
+    if moe_complementary_seq_aux_loss:
+        assert (
+            (moe_device_level_aux_loss_coeff is None) and
+            (moe_comm_aux_loss_coeff is None)
+            ), "moe_complementary_seq_aux_loss only used in deepseekV3, which means no other aux loss used"
+        probs_for_aux_loss = probs.view(seq_length, batch_size, -1)
+        sum_value = probs_for_aux_loss.sum(dim=-1, keepdim=True)
+        probs_for_aux_loss = probs_for_aux_loss / (sum_value + 1e-20)
+    seq_aux_loss = (cost_coeff * probs_for_aux_loss.mean(dim=0)).sum(dim=1).mean()
+    seq_aux_loss *= moe_aux_loss_coeff
+
+    if moe_device_level_aux_loss_coeff is not None:
+        num_group = (
+        parallel_state.get_expert_model_parallel_world_size()
+        )  # num_group equals to expert parallel size
+        device_aux_loss = (cost_coeff.view(batch_size, num_group, -1).mean(dim=2) *
+                           probs_for_aux_loss.mean(dim=0).view(batch_size, num_group, -1).sum(dim=2)).sum(dim=1).mean()
+        device_aux_loss *= moe_device_level_aux_loss_coeff
+        seq_aux_loss += device_aux_loss
+    if moe_comm_aux_loss_coeff is not None:
+        num_group = (
+        parallel_state.get_expert_model_parallel_world_size()
+        )  # num_group equals to expert parallel size
+        cost_coeff = routing_map.view(seq_length, batch_size, num_group, -1).any(dim=3).sum(dim=0).float()
+        cost_coeff.div_(seq_length *  moe_router_topk_limited_devices/ num_group)
+        comm_aux_loss = (cost_coeff *
+                           probs_for_aux_loss.mean(dim=0).view(batch_size, num_group, -1).sum(dim=2)).sum(dim=1).mean()
+        comm_aux_loss *= moe_comm_aux_loss_coeff
+        seq_aux_loss += comm_aux_loss
+
+    return seq_aux_loss
+
+def topk_softmax_with_capacity(
+    logits: torch.Tensor,
+    topk: int,
+    capacity_factor: Optional[float] = None,
+    pad_to_capacity: bool = False,
+    drop_policy: str = "probs",
+    use_pre_softmax: bool = False,
+    moe_router_topk_limited_devices: int = None,
+    moe_router_topk_scaling_factor: float = None,
+    device_level_capacity: Optional[bool] = False,
+    use_sigmoid: bool = False,
+    norm_topk_prob: bool = False,
+    num_node_group: int = None,
+    e_score_correction_bias: torch.Tensor = None,
+    deterministic_mode: bool = False,
+):
+    """Apply capacity and padding to the top-k selection.
+    Args:
+        logits (torch.Tensor): Logits tensor.
+        topk (int): The number of experts to select for each token.
+        capacity_factor (int): The capacity factor of each expert. Will drop tokens if the number
+                               of tokens exceeds the capacity.
+        pad_to_capacity (bool): Whether to need padding in token drop mode.
+        drop_policy (str): The policy to drop tokens. Can be either "prob" or "position".
+                           If "prob", the tokens with the lowest probabilities will be dropped.
+                           If "position", tokens at the end of each batch will be dropped.
+        use_pre_softmax (bool): Whether to apply softmax before top-k selection.
+        moe_router_topk_limited_devices (int): Number of expert parallel ranks to consider for
+            each token during routing. None means no device limitation.
+        moe_router_topk_scaling_factor (float): Scaling factor for routing score in top-k
+            selection, only works when use_pre_softmax enabled.
+        deterministic_mode (bool): Deprecated.
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            - routing_probs (torch.Tensor): A tensor of shape [num_tokens, num_experts] containing
+              the routing probabilities for each token to each expert.
+            - routing_map (torch.Tensor): A mask tensor of shape [num_tokens, num_experts]
+              indicating which experts were selected for each token. True values represent
+              the selected experts.
+            - tokens_per_expert (torch.Tensor): A tensor of shape [num_experts] containing
+              the number of local tokens assigned to each expert before dropping and padding.
+    """
+    assert logits.dim() == 2, f"Expected 2D logits [num_tokens, num_experts], got {logits.dim()}."
+    num_tokens = logits.shape[0]
+    num_experts = logits.shape[1]
+    if use_pre_softmax:
+        # Pre softmax
+        if use_sigmoid:
+            scores = torch.sigmoid(logits).type_as(logits)
+        else:
+            scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits)
+
+        if e_score_correction_bias is not None:
+            scores_gate = scores + e_score_correction_bias.unsqueeze(0)  #correction only used in router not in multiplied ffn output
+        else:
+            scores_gate = scores
+
+        if moe_router_topk_limited_devices:
+            if num_node_group:
+                top_indices = node_limited_topk(
+                    scores_gate, topk, num_tokens, num_experts, moe_router_topk_limited_devices, num_node_group
+                )
+                probs = scores.gather(1, top_indices)
+            else:
+                probs, top_indices = device_limited_topk(
+                    scores, topk, num_tokens, num_experts, moe_router_topk_limited_devices
+                )
+        else:
+            probs, top_indices = torch.topk(scores, k=topk, dim=1)
+
+        # Normalize the probs.
+        if norm_topk_prob:
+            assert use_sigmoid, f"norm_topk_prob only work with use_sigmoid=True, but get {use_sigmoid}"
+            denominator = probs.sum(dim=-1, keepdim=True) + 1e-20
+            probs = probs / denominator
+        if moe_router_topk_scaling_factor:
+            probs = probs * moe_router_topk_scaling_factor
+    else:
+        # Post softmax
+        if topk == 1:
+            # Requires applying softmax before selecting the top-k when k is 1,
+            # since softmax on a [num_tokens, 1] would yield a zero gradient.
+            raise ValueError("Please use --moe-router-pre-softmax when topk is 1.")
+        assert (
+            moe_router_topk_scaling_factor is None
+        ), "moe_router_topk_scaling_factor is not supported with post-softmax"
+        if moe_router_topk_limited_devices:
+            if num_node_group:
+                scores, top_indices = node_limited_topk(
+                    logits, topk, num_tokens, num_experts, moe_router_topk_limited_devices, num_node_group
+                )
+            else:
+                scores, top_indices = device_limited_topk(
+                    logits, topk, num_tokens, num_experts, moe_router_topk_limited_devices
+                )
+        else:
+            scores, top_indices = torch.topk(logits, k=topk, dim=1)
+        probs = torch.softmax(scores, dim=-1, dtype=torch.float32).type_as(logits)
+
+    # TODO Try using element-wise operations instead of scatter?
+    topk_masked_gates = torch.zeros_like(logits).scatter(1, top_indices, probs)
+    topk_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
+    tokens_per_expert = topk_map.sum(dim=0)
+
+    if capacity_factor is None:
+        # TopK without capacity
+        return topk_masked_gates, topk_map, tokens_per_expert
+    elif device_level_capacity:
+        assert drop_policy=='probs', f"only support 'probs' for device_level capacity, but get {drop_policy}"
+        num_group = (
+        parallel_state.get_expert_model_parallel_world_size()
+        )  # num_group equals to expert parallel size
+        device_expert_capacity = get_capacity(
+            num_tokens=num_tokens * topk, num_experts=num_experts, capacity_factor=capacity_factor
+        )*num_experts//num_group
+        # Maskout exceeded tokens
+        if drop_policy == "probs":
+            topk_masked_group_gates = topk_masked_gates.view(num_tokens, num_group, -1)
+            topk_masked_group_gates = topk_masked_group_gates.permute(0,2,1).reshape(-1, num_group)
+            _, capacity_indices = torch.topk(
+                topk_masked_group_gates, k=device_expert_capacity, dim=0, sorted=False
+            )
+            capacity_mask = torch.zeros([num_tokens*num_experts//num_group, num_group], device=logits.device).scatter(0, capacity_indices, 1).bool()
+            capacity_mask = capacity_mask.view(num_tokens, num_experts//num_group, num_group).permute(0,2,1).reshape(num_tokens, -1)
+        else:
+            raise ValueError(f"Invalid drop_policy: {drop_policy}")
+
+        if pad_to_capacity:
+            final_map = capacity_mask
+            final_probs = topk_masked_gates * final_map
+        else:
+            # Get exceed mask and maskout exceeded probs and indices
+            final_map = torch.logical_and(topk_map, capacity_mask)
+            final_probs = topk_masked_gates * final_map
+        return final_probs, final_map, tokens_per_expert
+    else:
+        # TopK with capacity
+        expert_capacity = get_capacity(
+            num_tokens=num_tokens * topk, num_experts=num_experts, capacity_factor=capacity_factor
+        )
+
+        # Maskout exceeded tokens
+        if drop_policy == "probs":
+            _, capacity_indices = torch.topk(
+                topk_masked_gates, k=expert_capacity, dim=0, sorted=False
+            )
+            capacity_mask = torch.zeros_like(logits).scatter(0, capacity_indices, 1).bool()
+        elif drop_policy == "position":
+            _, capacity_indices = torch.topk(topk_map.int(), k=expert_capacity, dim=0, sorted=False)
+            capacity_mask = torch.zeros_like(logits).scatter(0, capacity_indices, 1).bool()
+        else:
+            raise ValueError(f"Invalid drop_policy: {drop_policy}")
+
+        if pad_to_capacity:
+            final_map = capacity_mask
+            final_probs = topk_masked_gates * final_map
+        else:
+            # Get exceed mask and maskout exceeded probs and indices
+            final_map = torch.logical_and(topk_map, capacity_mask)
+            final_probs = topk_masked_gates * final_map
+        return final_probs, final_map, tokens_per_expert
+
+
+megatron.core.transformer.moe.moe_utils.sequence_load_balancing_loss_func = sequence_load_balancing_loss_func
+megatron.core.transformer.moe.moe_utils.topk_softmax_with_capacity = topk_softmax_with_capacity
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/cuda_patch/multi_latent_attention.py b/megatron-lm-musa-patch/cuda_patch/multi_latent_attention.py
new file mode 100644
index 000000000..098b284c2
--- /dev/null
+++ b/megatron-lm-musa-patch/cuda_patch/multi_latent_attention.py
@@ -0,0 +1,379 @@
+import math
+from dataclasses import dataclass
+from typing import Union
+
+import torch
+
+from megatron.core import parallel_state
+from megatron.core.models.common.embeddings import (
+    YarnRotaryEmbedding,
+    _yarn_get_mscale,
+    apply_rotary_pos_emb,
+)
+from megatron.core.transformer.attention import Attention
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import MLATransformerConfig
+
+from megatron.core.transformer.multi_latent_attention import MLASelfAttentionSubmodules
+
+class MultiLatentAttention(Attention):
+    """Multi-Latent Attention layer abstract class.
+
+    This layer only contains common modules required for the "self attn" and
+    "cross attn" specializations.
+    """
+
+    def __init__(
+        self,
+        config: MLATransformerConfig,
+        submodules: Union[MLASelfAttentionSubmodules],
+        layer_number: int,
+        attn_mask_type: AttnMaskType,
+        attention_type: str,
+        cp_comm_type: str = None,
+    ) -> None:
+        world_size = parallel_state.get_tensor_model_parallel_world_size()
+        assert (
+            world_size == 1
+        ), "MLA is not supported with Tensor Parallelism yet, \
+        use Expert Parallelism and Pipeline Parallelism for better performance."
+
+        super().__init__(
+            config=config,
+            submodules=submodules,
+            layer_number=layer_number,
+            attention_type=attention_type,
+            attn_mask_type=attn_mask_type,
+        )
+
+        self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads
+
+        self.q_head_dim = self.config.qk_head_dim + self.config.qk_pos_emb_head_dim
+
+        mscale = _yarn_get_mscale(self.config.rotary_scaling_factor, self.config.mscale)
+        self.softmax_scale = mscale * mscale / math.sqrt(self.q_head_dim)
+
+        self.rotary_pos_emb = YarnRotaryEmbedding(
+            self.config.qk_pos_emb_head_dim,
+            rotary_base=self.config.rotary_base,
+            scaling_factor=self.config.rotary_scaling_factor,
+            original_max_position_embeddings=self.config.max_position_embeddings,
+            beta_fast=self.config.beta_fast,
+            beta_slow=self.config.beta_slow,
+            mscale=self.config.mscale,
+            mscale_all_dim=self.config.mscale_all_dim,
+        )
+
+        self.core_attention = build_module(
+            submodules.core_attention,
+            config=self.config,
+            layer_number=self.layer_number,
+            attn_mask_type=self.attn_mask_type,
+            attention_type=self.attention_type,
+            softmax_scale=self.softmax_scale,
+            k_channels=self.q_head_dim,
+            v_channels=self.config.v_head_dim,
+            cp_comm_type=cp_comm_type,
+        )
+
+        # Output.
+        self.linear_proj = build_module(
+            submodules.linear_proj,
+            self.query_projection_size,
+            self.config.hidden_size,
+            config=self.config,
+            init_method=self.config.output_layer_init_method,
+            bias=self.config.add_bias_linear,
+            input_is_parallel=True,
+            skip_bias_add=True,
+            is_expert=False,
+            tp_comm_buffer_name='proj',
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        key_value_states=None,
+        inference_params=None,
+        rotary_pos_emb=None,
+        rotary_pos_cos=None,
+        rotary_pos_sin=None,
+        attention_bias=None,
+        packed_seq_params=None,
+        position_ids=None,
+        sequence_len_offset=None,
+    ):
+        """Forward pass for multi-latent attention"""
+        assert rotary_pos_emb is None, "Rotary position embeddings should not be passed into MLA."
+        assert attention_bias is None, "Attention bias should not be passed into MLA."
+        assert (
+            rotary_pos_cos is None and rotary_pos_sin is None
+        ), "MLA does not support Flash Decoding"
+
+        # hidden_states: [sq, b, h]
+
+        # =====================
+        # Query, Key, and Value
+        # =====================
+        # Get the query, key and value tensors based on the type of attention -
+        # self or cross attn.
+        # query: [96, 1, 16, 128], key:[96, 1, 16, 128], value:[96, 1, 16, 128]
+        query, key, value = self.get_query_key_value_tensors(
+            hidden_states,
+            key_value_states,
+            position_ids,
+            packed_seq_params,
+            inference_params=inference_params,
+        )
+
+        # ===================================================
+        # Adjust key, value for inference
+        # ===================================================
+        # rotary_pos_emb = None
+        query, key, value, _, attn_mask_type = self._adjust_key_value_for_inference(
+            inference_params, query, key, value, rotary_pos_emb=None
+        )
+
+        # ==================================
+        # core attention computation
+        # ==================================
+        # Need corresponding TE change
+        if self.checkpoint_core_attention and self.training:
+            core_attn_out = self._checkpointed_attention_forward(
+                query, key, value, attention_mask, packed_seq_params=packed_seq_params
+            )
+        else:
+            core_attn_out = self.core_attention(
+                query,
+                key,
+                value,
+                attention_mask,
+                packed_seq_params=packed_seq_params,
+                attn_mask_type=attn_mask_type,
+            )
+
+        if packed_seq_params is not None:
+            # reshape to same output shape as unpacked case
+            # (t, np, hn) -> (t, b=1, h=np*hn)
+            # t is the pack size = sum (sq_i)
+            # note that batch is a dummy dimension in the packed case
+            core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1)
+
+        # =================
+        # Output. [sq, b, h]
+        # =================
+        output, bias = self.linear_proj(core_attn_out)
+
+        return output, bias
+
+
+class MLASelfAttention(MultiLatentAttention):
+    """MLA Self-attention layer class
+
+    Self-attention layer takes input with size [s, b, h]
+    and returns output of the same size.
+    """
+
+    def __init__(
+        self,
+        config: MLATransformerConfig,
+        submodules: MLASelfAttentionSubmodules,
+        layer_number: int,
+        attn_mask_type=AttnMaskType.padding,
+        cp_comm_type: str = None,
+    ):
+        super().__init__(
+            config=config,
+            submodules=submodules,
+            layer_number=layer_number,
+            attn_mask_type=attn_mask_type,
+            attention_type="self",
+        )
+
+        if self.config.q_lora_rank is None:
+            # Not projectiing query
+            self.linear_q_proj = build_module(
+                submodules.linear_q_proj,
+                self.config.hidden_size,
+                self.config.num_attention_heads * self.q_head_dim,
+                config=self.config,
+                init_method=self.config.init_method,
+                gather_output=False,
+                bias=False,
+                skip_bias_add=False,
+                is_expert=False,
+            )
+
+        else:
+
+            self.linear_q_down_proj = build_module(
+                submodules.linear_q_down_proj,
+                self.config.hidden_size,
+                self.config.q_lora_rank,
+                config=self.config,
+                init_method=self.config.init_method,
+                gather_output=False,
+                bias=False,
+                skip_bias_add=False,
+                is_expert=False,
+            )
+
+            self.linear_q_up_proj = build_module(
+                submodules.linear_q_up_proj,
+                self.config.q_lora_rank,
+                self.config.num_attention_heads * self.q_head_dim,
+                config=self.config,
+                init_method=self.config.init_method,
+                gather_output=False,
+                bias=False,
+                skip_bias_add=False,
+                is_expert=False,
+            )
+
+        self.linear_kv_down_proj = build_module(
+            submodules.linear_kv_down_proj,
+            self.config.hidden_size,
+            self.config.kv_lora_rank + self.config.qk_pos_emb_head_dim,
+            config=self.config,
+            init_method=self.config.init_method,
+            gather_output=False,
+            bias=False,
+            skip_bias_add=False,
+            is_expert=False,
+        )
+
+        self.linear_kv_up_proj = build_module(
+            submodules.linear_kv_up_proj,
+            self.config.kv_lora_rank,
+            self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim),
+            config=self.config,
+            init_method=self.config.init_method,
+            gather_output=False,
+            bias=False,
+            skip_bias_add=False,
+            is_expert=False,
+        )
+
+        if self.config.q_lora_rank is not None:
+            self.q_layernorm = build_module(
+                submodules.q_layernorm,
+                hidden_size=self.config.q_lora_rank,
+                config=self.config,
+                eps=self.config.layernorm_epsilon,
+            )
+
+        self.kv_layernorm = build_module(
+            submodules.kv_layernorm,
+            hidden_size=self.config.kv_lora_rank,
+            config=self.config,
+            eps=self.config.layernorm_epsilon,
+        )
+
+    def get_query_key_value_tensors(
+        self,
+        hidden_states,
+        key_value_states=None,
+        position_ids=None,
+        packed_seq_params=None,
+        inference_params=None,
+    ):
+        """
+        Derives `query`, `key` and `value` tensors from `hidden_states`.
+        """
+        # s = sequence length, b = batch size, h = hidden size, n = num attention heads
+        # Attention heads [s, b, n*h]
+        assert (
+            hidden_states.ndim == 3
+        ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D"
+        q_len, bsz, _ = hidden_states.size()
+
+        if self.config.q_lora_rank is not None:
+            q_compressed, _ = self.linear_q_down_proj(hidden_states)
+            q_compressed = self.q_layernorm(q_compressed)
+            q, _ = self.linear_q_up_proj(q_compressed)
+        else:
+            # hidden_states:[s, b, 2048], q: [s, b, n * 192]
+            q, _ = self.linear_q_proj(hidden_states)
+
+        # q: [s, b, n, 192]
+        q = q.view(q_len, bsz, self.num_attention_heads_per_partition, self.q_head_dim)
+
+        # q: [s, b, n, 128], q_pos_emb: [s, b, n, 64]
+        q_no_pe, q_pos_emb = torch.split(
+            q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1
+        )
+
+        # kv_combined: [s, b, 576]
+        kv_combined, _ = self.linear_kv_down_proj(hidden_states)
+
+        # kv_compressed:[s, b, 512], k_pos_emb: [s, b, 64]
+        kv_compressed, k_pos_emb = torch.split(
+            kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1
+        )
+
+        # kv: [s, b, 2048]
+        kv, _ = self.linear_kv_up_proj(self.kv_layernorm(kv_compressed))
+
+        # kv: [s, b, n, 256]
+        kv = kv.view(
+            q_len,
+            bsz,
+            self.num_attention_heads_per_partition,
+            self.config.qk_head_dim + self.config.v_head_dim,
+        )
+
+        # k_no_pe: [s, b, n, 128], value: [s, b, n, 128]
+        k_no_pe, value = torch.split(kv, [self.config.qk_head_dim, self.config.v_head_dim], dim=-1)
+
+        # rotary_pos_emb:[s, b, 1, 64]
+        rotary_pos_emb = self.rotary_pos_emb(max_seq_len=self.config.max_position_embeddings)
+
+        if len(rotary_pos_emb) == 2:
+            mscale = rotary_pos_emb[1]
+            rotary_pos_emb = rotary_pos_emb[0]
+
+        if inference_params is not None:
+            # add offset to the sequence start for inference
+            sequence_start = inference_params.sequence_len_offset
+            sequence_end = sequence_start + q_len
+            rotary_pos_emb = rotary_pos_emb[sequence_start:sequence_end]
+
+        # [s, b, 64] -> [s, b, 1, 64]
+        k_pos_emb = torch.unsqueeze(k_pos_emb, 2)
+
+        if packed_seq_params is not None:
+            cu_seqlens_q = packed_seq_params.cu_seqlens_q
+            cu_seqlens_kv = packed_seq_params.cu_seqlens_kv
+        else:
+            cu_seqlens_q = cu_seqlens_kv = None
+
+        # q_pos_emb: [s, b, n, 64], k_pos_emb:[s, b, 1, 64]
+        q_pos_emb = apply_rotary_pos_emb(
+            q_pos_emb, rotary_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q, mscale=mscale
+        )
+        k_pos_emb = apply_rotary_pos_emb(
+            k_pos_emb, rotary_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv, mscale=mscale
+        )
+
+        # query: [s, b, n, 192]
+        query = torch.cat([q_no_pe, q_pos_emb], dim=-1)
+
+        # key: [s, b, n, 192]
+        k_pos_emb = k_pos_emb.expand(-1, -1, self.config.num_attention_heads, -1)
+        key = torch.cat([k_no_pe, k_pos_emb], dim=-1)
+
+        query = query.contiguous()
+        key = key.contiguous()
+        value = value.contiguous()
+
+        return query, key, value
+
+
+import sys
+for k in sys.modules:
+    if k.startswith('megatron'):
+        for target in ['MLASelfAttention']:
+            if getattr(sys.modules[k], target, None):
+                setattr(sys.modules[k], target, MLASelfAttention)
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/cuda_patch/profiling.py b/megatron-lm-musa-patch/cuda_patch/profiling.py
new file mode 100644
index 000000000..17f2abe62
--- /dev/null
+++ b/megatron-lm-musa-patch/cuda_patch/profiling.py
@@ -0,0 +1,141 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import contextlib
+import os
+import pickle
+import time
+from datetime import datetime
+
+import torch
+
+# the number of warmup steps before the active step in each profiling cycle
+profile_freq = 4
+# how much memory allocation/free ops to record in memory snapshots
+MEMORY_SNAPSHOT_MAX_ENTRIES = 100000
+
+
+@contextlib.contextmanager
+def maybe_enable_profiling(args, global_step):
+    # get user defined profiler settings
+    enable_profiling = int(os.getenv("ENABLE_PROFILER", 0))
+     # fetch profiler related env
+    wait_steps = int(os.getenv("PROFILER_WAIT_STEPS", 0))
+    warmup_steps = int(os.getenv("PROFILER_WARMUP_STEPS", 3))
+    active_steps = int(os.getenv("PROFILER_ACTIVE_STEPS", 1))
+    repeat_num = int(os.getenv("PROFILER_REPEAT_NUM", 0))
+    profile_freq = int(os.getenv("PROFILER_FREQ", 1))
+    current_time = datetime.now().strftime("%Y.%m.%d-%H:%M:%S")
+    save_dir = os.getenv("PROFILER_SAVE_DIR", f"./profiler_result/{current_time}")
+    worker_name = os.getenv(
+        "PROFILER_WORKER_NAME", "rank" + str(torch.distributed.get_rank())
+    )
+    record_shapes = int(os.getenv("PROFILER_RECORD_SHAPES", 1))
+    profile_memory = int(os.getenv("PROFILER_PROFILE_MEMORY", 0))
+    with_stack = int(os.getenv("PROFILER_WITH_STACK", 1))
+    with_modules = int(os.getenv("PROFILER_WITH_MODULES", 1))
+    kineto_log_level = int(os.getenv("KINETO_LOG_LEVEL", 0))
+
+    if enable_profiling:
+        profile_freq = profile_freq
+
+        rank = torch.distributed.get_rank()
+
+        def trace_handler(prof):
+            curr_trace_dir_name = "iteration_" + str(prof.step_num)
+            curr_trace_dir = os.path.join(save_dir, curr_trace_dir_name)
+            if not os.path.exists(curr_trace_dir):
+                os.makedirs(curr_trace_dir, exist_ok=True)
+
+            print(f"Dumping profiler traces at step {prof.step_num}")
+            begin = time.monotonic()
+            prof.export_chrome_trace(f"{curr_trace_dir}/rank{rank}_trace.pt.trace.json")
+            print(
+                f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"
+            )
+
+        print(f"Profiling active. Traces will be saved at {save_dir}")
+
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir, exist_ok=True)
+
+        # warmup, active = WARMUP, 1
+        wait = profile_freq - (active_steps + warmup_steps)
+        assert (
+            wait >= 0
+        ), "profile_freq must be greater than or equal to warmup + active"
+        with torch.profiler.profile(
+            activities=[
+                torch.profiler.ProfilerActivity.CPU,
+                #torch.profiler.ProfilerActivity.MUSA,
+                torch.profiler.ProfilerActivity.CUDA,
+            ],
+            schedule=torch.profiler.schedule(wait=wait, warmup=warmup_steps, active=active_steps),
+            on_trace_ready=trace_handler,
+            record_shapes=record_shapes,
+            profile_memory=profile_memory,
+            with_stack=with_stack,
+            with_modules=with_modules,
+            start_step=global_step+1,
+        ) as torch_profiler:
+            yield torch_profiler
+    else:
+        torch_profiler = contextlib.nullcontext()
+        yield None
+
+
+@contextlib.contextmanager
+def maybe_enable_memory_snapshot(args, global_step: int = 0):
+    pass
+    # enable_snapshot = config.profiling.enable_memory_snapshot
+    # if enable_snapshot:
+    #     snapshot_folder = config.profiling.save_memory_snapshot_folder
+    #     snapshot_dir = os.path.join(config.job.dump_folder, snapshot_folder)
+    #     if not os.path.exists(snapshot_dir):
+    #         os.makedirs(snapshot_dir, exist_ok=True)
+    #     rank = torch.distributed.get_rank()
+
+    #     class MemoryProfiler:
+    #         def __init__(self, step_num: int, freq: int):
+    #             torch.musa.memory._record_memory_history(
+    #                 max_entries=MEMORY_SNAPSHOT_MAX_ENTRIES
+    #             )
+    #             # when resume training, we start from the last step
+    #             self.step_num = step_num
+    #             self.freq = freq
+
+    #         def step(self, exit_ctx: bool = False):
+    #             self.step_num += 1
+    #             if not exit_ctx and self.step_num % self.freq != 0:
+    #                 return
+    #             if not exit_ctx:
+    #                 curr_step = self.step_num
+    #                 dir_name = f"iteration_{curr_step}"
+    #             else:
+    #                 # dump as iteration_0_exit if OOM at iter 1
+    #                 curr_step = self.step_num - 1
+    #                 dir_name = f"iteration_{curr_step}_exit"
+    #             curr_snapshot_dir = os.path.join(snapshot_dir, dir_name)
+    #             if not os.path.exists(curr_snapshot_dir):
+    #                 os.makedirs(curr_snapshot_dir, exist_ok=True)
+    #             logger.info(f"Dumping memory snapshot at step {curr_step}")
+    #             begin = time.monotonic()
+    #             with open(
+    #                 f"{curr_snapshot_dir}/rank{rank}_memory_snapshot.pickle", "wb"
+    #             ) as output:
+    #                 pickle.dump(torch.musa.memory._snapshot(), output)
+    #             logger.info(
+    #                 f"Finished dumping memory snapshot in {time.monotonic() - begin:.2f} seconds"
+    #             )
+
+    #     logger.info(f"Memory profiler active. Snapshot will be saved at {snapshot_dir}")
+    #     profiler = MemoryProfiler(global_step, config.profiling.profile_freq)
+    #     try:
+    #         yield profiler
+    #     except torch.OutOfMemoryError as e:
+    #         profiler.step(exit_ctx=True)
+    # else:
+    #     yield None
diff --git a/megatron-lm-musa-patch/cuda_patch/profiling.py. and b/megatron-lm-musa-patch/cuda_patch/profiling.py. and
new file mode 100644
index 000000000..8d9eb5c72
--- /dev/null
+++ b/megatron-lm-musa-patch/cuda_patch/profiling.py. and	
@@ -0,0 +1,140 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import contextlib
+import os
+import pickle
+import time
+from datetime import datetime
+
+import torch
+
+# the number of warmup steps before the active step in each profiling cycle
+profile_freq = 4
+# how much memory allocation/free ops to record in memory snapshots
+MEMORY_SNAPSHOT_MAX_ENTRIES = 100000
+
+
+@contextlib.contextmanager
+def maybe_enable_profiling(args, global_step):
+    # get user defined profiler settings
+    enable_profiling = int(os.getenv("ENABLE_PROFILER", 1))
+     # fetch profiler related env
+    wait_steps = int(os.getenv("PROFILER_WAIT_STEPS", 0))
+    warmup_steps = int(os.getenv("PROFILER_WARMUP_STEPS", 3))
+    active_steps = int(os.getenv("PROFILER_ACTIVE_STEPS", 1))
+    repeat_num = int(os.getenv("PROFILER_REPEAT_NUM", 0))
+    profile_freq = int(os.getenv("PROFILER_FREQ", 0))
+    current_time = datetime.now().strftime("%Y.%m.%d-%H:%M:%S")
+    save_dir = os.getenv("PROFILER_SAVE_DIR", f"./profiler_result/{current_time}")
+    worker_name = os.getenv(
+        "PROFILER_WORKER_NAME", "rank" + str(torch.distributed.get_rank())
+    )
+    record_shapes = int(os.getenv("PROFILER_RECORD_SHAPES", 1))
+    profile_memory = int(os.getenv("PROFILER_PROFILE_MEMORY", 0))
+    with_stack = int(os.getenv("PROFILER_WITH_STACK", 1))
+    with_modules = int(os.getenv("PROFILER_WITH_MODULES", 1))
+    kineto_log_level = int(os.getenv("KINETO_LOG_LEVEL", 0))
+
+    if enable_profiling:
+        profile_freq = profile_freq
+
+        rank = torch.distributed.get_rank()
+
+        def trace_handler(prof):
+            curr_trace_dir_name = "iteration_" + str(prof.step_num)
+            curr_trace_dir = os.path.join(save_dir, curr_trace_dir_name)
+            if not os.path.exists(curr_trace_dir):
+                os.makedirs(curr_trace_dir, exist_ok=True)
+
+            print(f"Dumping profiler traces at step {prof.step_num}")
+            begin = time.monotonic()
+            prof.export_chrome_trace(f"{curr_trace_dir}/rank{rank}_trace.json")
+            print(
+                f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"
+            )
+
+        print(f"Profiling active. Traces will be saved at {save_dir}")
+
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir, exist_ok=True)
+
+        # warmup, active = WARMUP, 1
+        wait = profile_freq - (active_steps + warmup_steps)
+        assert (
+            wait >= 0
+        ), "profile_freq must be greater than or equal to warmup + active"
+        with torch.profiler.profile(
+            activities=[
+                torch.profiler.ProfilerActivity.CPU,
+                torch.profiler.ProfilerActivity.CUDA,
+            ],
+            schedule=torch.profiler.schedule(wait=wait, warmup=warmup_steps, active=active_steps),
+            on_trace_ready=trace_handler,
+            record_shapes=record_shapes,
+            profile_memory=profile_memory,
+            with_stack=with_stack,
+            with_modules=with_modules,
+        ) as torch_profiler:
+            torch_profiler.step_num = global_step
+            yield torch_profiler
+    else:
+        torch_profiler = contextlib.nullcontext()
+        yield None
+
+
+@contextlib.contextmanager
+def maybe_enable_memory_snapshot(args, global_step: int = 0):
+    pass
+    # enable_snapshot = config.profiling.enable_memory_snapshot
+    # if enable_snapshot:
+    #     snapshot_folder = config.profiling.save_memory_snapshot_folder
+    #     snapshot_dir = os.path.join(config.job.dump_folder, snapshot_folder)
+    #     if not os.path.exists(snapshot_dir):
+    #         os.makedirs(snapshot_dir, exist_ok=True)
+    #     rank = torch.distributed.get_rank()
+
+    #     class MemoryProfiler:
+    #         def __init__(self, step_num: int, freq: int):
+    #             torch.musa.memory._record_memory_history(
+    #                 max_entries=MEMORY_SNAPSHOT_MAX_ENTRIES
+    #             )
+    #             # when resume training, we start from the last step
+    #             self.step_num = step_num
+    #             self.freq = freq
+
+    #         def step(self, exit_ctx: bool = False):
+    #             self.step_num += 1
+    #             if not exit_ctx and self.step_num % self.freq != 0:
+    #                 return
+    #             if not exit_ctx:
+    #                 curr_step = self.step_num
+    #                 dir_name = f"iteration_{curr_step}"
+    #             else:
+    #                 # dump as iteration_0_exit if OOM at iter 1
+    #                 curr_step = self.step_num - 1
+    #                 dir_name = f"iteration_{curr_step}_exit"
+    #             curr_snapshot_dir = os.path.join(snapshot_dir, dir_name)
+    #             if not os.path.exists(curr_snapshot_dir):
+    #                 os.makedirs(curr_snapshot_dir, exist_ok=True)
+    #             logger.info(f"Dumping memory snapshot at step {curr_step}")
+    #             begin = time.monotonic()
+    #             with open(
+    #                 f"{curr_snapshot_dir}/rank{rank}_memory_snapshot.pickle", "wb"
+    #             ) as output:
+    #                 pickle.dump(torch.musa.memory._snapshot(), output)
+    #             logger.info(
+    #                 f"Finished dumping memory snapshot in {time.monotonic() - begin:.2f} seconds"
+    #             )
+
+    #     logger.info(f"Memory profiler active. Snapshot will be saved at {snapshot_dir}")
+    #     profiler = MemoryProfiler(global_step, config.profiling.profile_freq)
+    #     try:
+    #         yield profiler
+    #     except torch.OutOfMemoryError as e:
+    #         profiler.step(exit_ctx=True)
+    # else:
+    #     yield None
diff --git a/megatron-lm-musa-patch/cuda_patch/router.py b/megatron-lm-musa-patch/cuda_patch/router.py
new file mode 100644
index 000000000..57f01ee73
--- /dev/null
+++ b/megatron-lm-musa-patch/cuda_patch/router.py
@@ -0,0 +1,97 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Callable
+import torch
+
+from megatron.core import parallel_state
+from megatron.core.transformer.moe.moe_utils import (
+    MoEAuxLossAutoScaler,
+    save_to_aux_losses_tracker,
+)
+from megatron.core.transformer.moe.router import TopKRouter
+
+from .moe_utils import (
+    sequence_load_balancing_loss_func,
+    topk_softmax_with_capacity,
+)
+
+
+def init(self, config):
+    """Initialize the zero token dropping router.
+
+    Args:
+        config (TransformerConfig): The configuration for the transformer model.
+    """
+    super(TopKRouter, self).__init__(config=config)
+    self.topk = self.config.moe_router_topk
+    self.routing_type = self.config.moe_router_load_balancing_type
+    self.input_jitter = None
+    self.e_score_correction_bias = None
+    self.moe_noaux_gamma = self.config.moe_noaux_gamma
+    if self.moe_noaux_gamma:
+        self.tokens_per_expert = torch.nn.Parameter(
+             torch.empty((self.config.num_moe_experts)), requires_grad=False
+        )
+        self.e_score_correction_bias = torch.nn.Parameter(
+             torch.empty((self.config.num_moe_experts)), requires_grad=False
+        )
+        if config.perform_initialization:
+            # initialize bias to zero.
+            with torch.no_grad():
+                self.e_score_correction_bias.zero_()
+                self.tokens_per_expert.zero_()
+
+def seq_aux_loss_load_balancing(self, logits: torch.Tensor, bsz: int, seq_length: int):
+    """Apply loss-based load balancing to the logits tensor."""
+
+    probs, routing_map, tokens_per_expert = topk_softmax_with_capacity(
+        logits,
+        self.topk,
+        capacity_factor=self.config.moe_expert_capacity_factor,
+        pad_to_capacity=self.config.moe_pad_expert_input_to_capacity,
+        drop_policy=self.config.moe_token_drop_policy,
+        use_pre_softmax=self.config.moe_router_pre_softmax,
+        moe_router_topk_limited_devices=self.config.moe_router_topk_limited_devices,
+        moe_router_topk_scaling_factor=self.config.moe_router_topk_scaling_factor,
+        use_sigmoid=self.config.moe_router_use_sigmoid,
+        norm_topk_prob=self.config.moe_router_norm_topk_prob,
+        deterministic_mode=self.config.deterministic_mode,
+        device_level_capacity=self.config.moe_device_level_capacity,
+        num_node_group=self.config.moe_router_num_node_group,
+        e_score_correction_bias=self.e_score_correction_bias,
+    )
+
+    if self.training:
+        if self.config.moe_router_use_sigmoid == "sigmoid":
+            scores = torch.sigmoid(logits)
+        else:
+            scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
+        aux_loss_func = partial(
+            sequence_load_balancing_loss_func,
+            probs=scores,
+            routing_map=routing_map,
+            tokens_per_expert=tokens_per_expert,
+            batch_size=bsz,
+            seq_length=seq_length,
+            topk=self.topk,
+            moe_router_topk_limited_devices=self.config.moe_router_topk_limited_devices,
+            moe_device_level_aux_loss_coeff=self.config.moe_device_level_aux_loss_coeff,
+            moe_comm_aux_loss_coeff=self.config.moe_comm_aux_loss_coeff,
+            moe_complementary_seq_aux_loss=self.config.moe_complementary_seq_aux_loss,
+        )
+        probs = self.apply_load_balancing_loss(
+            activation=probs, load_balancing_loss_func=aux_loss_func
+        )
+
+        if self.moe_noaux_gamma and torch.is_grad_enabled():
+            with torch.no_grad():
+                mean_token = logits.shape[0] * self.topk / logits.shape[1]
+                self.tokens_per_expert += (tokens_per_expert /mean_token).to(self.tokens_per_expert.device)
+    return probs, routing_map
+
+
+import megatron.core.transformer.moe.router
+megatron.core.transformer.moe.router.TopKRouter.seq_aux_loss_load_balancing = seq_aux_loss_load_balancing
+megatron.core.transformer.moe.router.TopKRouter.__init__ = init
diff --git a/megatron-lm-musa-patch/cuda_patch/theoretical_memory_usage.py b/megatron-lm-musa-patch/cuda_patch/theoretical_memory_usage.py
new file mode 100644
index 000000000..2f977bd18
--- /dev/null
+++ b/megatron-lm-musa-patch/cuda_patch/theoretical_memory_usage.py
@@ -0,0 +1,220 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+"""Computes theoretical memory footprint for model training."""
+
+
+import math
+
+NUM_BYTES_IN_MEGABYTE = 1024 * 1024
+
+def compute_weight_and_optimizer_memory(args, verbose=False):
+    # Attention projection size.
+    attn_dim = args.kv_channels
+    kv_projection_size = attn_dim * args.num_attention_heads
+    query_projection_size = attn_dim * args.num_attention_heads
+    ## MLA
+    if args.kv_lora_rank:
+        kv_projection_size = attn_dim * args.num_attention_heads + args.kv_lora_rank
+    if args.q_lora_rank:
+        query_projection_size = attn_dim * args.num_attention_heads + args.q_lora_rank
+
+    output_projection_size = attn_dim * args.num_attention_heads
+    ## Group Query Attention.
+    if args.group_query_attention:
+        kv_projection_size = args.num_query_groups / args.num_attention_heads * kv_projection_size
+    else:
+        attn_size = 1
+    attn_size = query_projection_size + 2 * kv_projection_size + output_projection_size
+    attn_multiplier = attn_size / 2 / args.hidden_size
+
+    # swiglu
+    gated_linear_multiplier = 3 / 2 if args.swiglu else 1
+
+    # MoE or Dense
+    num_experts = 0 if args.num_experts is None else args.num_experts
+    shared_expert_ffn_hidden_size = 0 if args.moe_shared_expert_intermediate_size else args.moe_shared_expert_intermediate_size
+    num_shared_expert = args.moe_shared_expert_intermediate_size // args.moe_ffn_hidden_size
+    num_moe_layer = 0 if args.moe_layer_freq is None else len(args.moe_layer_freq)
+    num_dense_layer = args.num_layers - num_moe_layer
+    mlp_multiplier_dense = num_dense_layer * args.ffn_hidden_size
+    mlp_multiplier_moe =  num_moe_layer * num_experts * args.moe_ffn_hidden_size
+    mlp_multiplier_shared_expert = num_moe_layer * num_shared_expert * shared_expert_ffn_hidden_size
+    mlp_multiplier = (mlp_multiplier_dense + mlp_multiplier_moe + mlp_multiplier_shared_expert) /args.num_layers/args.hidden_size
+
+    num_parameters_in_transformer_layers = (
+        2
+        * args.num_layers
+        * args.hidden_size
+        * args.hidden_size
+        * (
+            # Attention.
+            attn_multiplier
+            # MLP.
+            # + ((args.ffn_hidden_size / args.hidden_size) * num_experts * gated_linear_multiplier)
+            + mlp_multiplier * gated_linear_multiplier
+            # Router
+            +  num_experts / args.hidden_size * 2
+            # Transformer layernorms.
+            + (2 / args.hidden_size)
+            # Final layernorm.
+            + (1 / (args.num_layers * args.hidden_size))
+        )
+    )
+
+    embedding_size = args.hidden_size * args.padded_vocab_size
+    if args.untie_embeddings_and_output_weights:
+        num_parameters_in_embedding_layers = 2 * embedding_size
+    else:
+        num_parameters_in_embedding_layers = embedding_size
+
+    # TODO; add MTP block and projection
+    num_parameters_in_mtp = 0
+
+    num_total_parameters = num_parameters_in_transformer_layers + num_parameters_in_embedding_layers + num_parameters_in_mtp
+    if verbose:
+        print(
+            f"Number of parameters in transformer layers in billions: "
+            f"{num_parameters_in_transformer_layers / 10**9: .2f}"
+        )
+        print(
+            f"Number of parameters in embedding layers in billions: "
+            f"{num_parameters_in_embedding_layers / 10**9:.2f}"
+        )
+        print(f"Total number of parameters in billions: {num_total_parameters / 10**9:.2f}")
+
+    # Most loaded model shard has (1/pp_size transformer layers + 1 embedding layer) / tp_size.
+    num_parameters_on_most_loaded_model_shard = (
+        (num_parameters_in_transformer_layers / args.pipeline_model_parallel_size) + embedding_size
+    ) / args.tensor_model_parallel_size
+    if args.untie_embeddings_and_output_weights and args.pipeline_model_parallel_size == 1:
+        num_parameters_on_most_loaded_model_shard += (
+            embedding_size / args.tensor_model_parallel_size
+        )
+    if verbose:
+        print(
+            f"Number of parameters in most loaded shard in billions: "
+            f"{num_parameters_on_most_loaded_model_shard / 10**9:.4f}"
+        )
+
+    if args.pipeline_model_parallel_size > 1:
+        # Other shards just have (1/pp_size transformer layers) / tp_size.
+        num_parameters_on_other_model_shards = num_parameters_in_transformer_layers / (
+            args.pipeline_model_parallel_size * args.tensor_model_parallel_size
+        )
+        if verbose:
+            print(
+                f"Number of parameters in other shards in billions: "
+                f"{num_parameters_on_other_model_shards / 10**9:.4f}"
+            )
+
+    num_bytes_per_parameter = (
+        18 if not args.use_distributed_optimizer else 6 + (12 / args.data_parallel_size)
+    )
+    weight_and_optimizer_memory = (
+        num_parameters_on_most_loaded_model_shard * num_bytes_per_parameter
+    )
+
+    return weight_and_optimizer_memory
+
+
+def compute_activation_memory(args, num_microbatches, verbose=False):
+    # Using formula in Table 2 of https://arxiv.org/pdf/2205.05198.pdf.
+    # We are trying to compute the maximum activation footprint, so all calculations in this
+    # function are for the first pipeline stage.
+
+    # TODO: This function needs to take into account query_projection_size potentially being
+    # different from hidden_size.
+
+    # Memory footprint from transformer layer (self-attention and MLP).
+    activation_memory = (args.seq_length * args.micro_batch_size * args.hidden_size) * (
+        18 + (4 * (args.ffn_hidden_size / args.hidden_size))
+    )
+    if verbose:
+        print(
+            f"Activation memory footprint per transformer layer: "
+            f"{activation_memory / NUM_BYTES_IN_MEGABYTE / args.tensor_model_parallel_size:.1f} MB"
+        )
+    activation_memory *= args.num_layers
+
+    # Now add activation memory required for input embeddings, last LayerNorm and output layer.
+
+    # Input to embedding (pp_size microbatches in flight).
+    activation_memory += (
+        8 * args.seq_length * args.micro_batch_size * args.pipeline_model_parallel_size
+    )
+    # Dropout in embedding layer (pp_size microbatches in flight).
+    activation_memory += (
+        args.seq_length
+        * args.micro_batch_size
+        * args.hidden_size
+        * args.pipeline_model_parallel_size
+    )
+
+    # Multiply by interleaved PP memory factor.
+    if args.virtual_pipeline_model_parallel_size is not None:
+        interleaved_schedule_memory_penalty = 1 + (
+            (args.pipeline_model_parallel_size - 1)
+            / (args.pipeline_model_parallel_size * args.virtual_pipeline_model_parallel_size)
+        )
+        in_flight_microbatches = math.ceil(
+            interleaved_schedule_memory_penalty * args.pipeline_model_parallel_size
+        )
+        if verbose:
+            print(
+                f"Memory penalty from interleaved schedule: {interleaved_schedule_memory_penalty:.2f}"
+            )
+            print(f"Number of in-flight microbatches: {in_flight_microbatches}")
+        activation_memory *= interleaved_schedule_memory_penalty
+
+    # If using non-interleaved schedule, number of microbatches in pipeline can be less than pp_size,
+    # so discount accordingly.
+    if args.virtual_pipeline_model_parallel_size is None and args.pipeline_model_parallel_size > 1:
+        if num_microbatches is not None:
+            activation_memory *= min(1, num_microbatches / args.pipeline_model_parallel_size)
+            in_flight_microbatches = min(num_microbatches, args.pipeline_model_parallel_size)
+        else:
+            in_flight_microbatches = args.pipeline_model_parallel_size
+        if verbose:
+            print(f"Number of in-flight microbatches: {in_flight_microbatches}")
+
+    if args.pipeline_model_parallel_size == 1:
+        # Inputs to output layer and CE loss.
+        activation_memory += (
+            args.seq_length
+            * args.micro_batch_size
+            * args.hidden_size
+            * 4
+            * (1 + (args.padded_vocab_size / args.hidden_size))
+        )
+
+    # Activation memory is partitioned by TP size due to tensor and sequence model parallelism.
+    return activation_memory / args.tensor_model_parallel_size
+
+
+def report_theoretical_memory(args, num_microbatches=None, verbose=False):
+    weight_and_optimizer_memory = (
+        compute_weight_and_optimizer_memory(args, verbose=verbose) / NUM_BYTES_IN_MEGABYTE
+    )
+
+    # Formulae here assume sequence parallelism and selective activation recomputation.
+    if not args.sequence_parallel or args.recompute_granularity != 'selective':
+        print(
+            f"Theoretical memory footprints: weight and optimizer={weight_and_optimizer_memory:.2f} MB"
+        )
+        return
+
+    activation_memory = (
+        compute_activation_memory(args, num_microbatches=num_microbatches, verbose=verbose)
+        / NUM_BYTES_IN_MEGABYTE
+    )
+    total_memory = weight_and_optimizer_memory + activation_memory
+
+    print(
+        f"Theoretical memory footprints: weight and optimizer={weight_and_optimizer_memory:.2f} MB, "
+        f"activation={activation_memory:.2f} MB, total={total_memory:.2f} MB\n"
+    )
+
+import megatron.training.theoretical_memory_usage
+megatron.training.theoretical_memory_usage.compute_weight_and_optimizer_memory = compute_weight_and_optimizer_memory
+megatron.training.theoretical_memory_usage.report_theoretical_memory = report_theoretical_memory
+megatron.training.theoretical_memory_usage.compute_activation_memory = compute_activation_memory
diff --git a/megatron-lm-musa-patch/cuda_patch/training.py b/megatron-lm-musa-patch/cuda_patch/training.py
new file mode 100644
index 000000000..ee3a83e6b
--- /dev/null
+++ b/megatron-lm-musa-patch/cuda_patch/training.py
@@ -0,0 +1,813 @@
+from datetime import datetime
+
+import gc
+import os
+import sys
+import time
+import torch
+from megatron.core import mpu
+from megatron.core.transformer.moe.moe_utils import track_moe_metrics
+from megatron.core.transformer.moe.router import TopKRouter
+from megatron.training.global_vars import (
+    get_args,
+    get_timers,
+    get_tensorboard_writer,
+    get_wandb_writer,
+    get_one_logger
+)
+    # get_num_microbatches
+from megatron.core.num_microbatches_calculator import get_num_microbatches
+from megatron.training.utils import (
+    report_memory,
+    print_rank_last
+)
+from megatron.core.utils import (
+    check_param_hashes_across_dp_replicas
+)
+from megatron.training.theoretical_memory_usage import report_theoretical_memory
+from megatron.training import one_logger_utils
+from megatron.training.initialize import write_args_to_tensorboard
+from megatron.core.distributed import finalize_model_grads
+from megatron.core.distributed import DistributedDataParallel as DDP
+from megatron.training.training import (
+    print_datetime,
+    save_checkpoint_and_time,
+    train_step,
+    evaluate_and_print_results,
+    _TRAIN_START_TIME
+)
+from megatron.training.async_utils import maybe_finalize_async_save
+from megatron.core.num_microbatches_calculator import (
+    get_current_global_batch_size,
+    get_current_running_global_batch_size,
+    get_num_microbatches,
+    update_num_microbatches
+)
+from megatron.training.utils import (
+    calc_params_l2_norm,
+    check_adlr_autoresume_termination,
+    print_rank_0,
+    print_rank_last,
+    report_memory
+)
+from megatron.training import ft_integration
+from megatron.training.global_vars import (
+    get_args,
+    get_signal_handler,
+    get_timers,
+    get_tensorboard_writer,
+    get_wandb_writer,
+    get_one_logger
+)
+from .profiling import (
+    maybe_enable_profiling,
+    maybe_enable_memory_snapshot
+)
+from megatron.training.training import (
+    enable_forward_pre_hook,
+    disable_forward_pre_hook,
+    post_training_step_callbacks,
+    checkpoint_and_decide_exit
+)
+
+
+def throughput_calculator(args, elapsed_time_per_iter, consumed_tokens_per_iter):
+    # training_time = elapsed_time
+    system_throughput = float(consumed_tokens_per_iter) / elapsed_time_per_iter
+    world_size = args.world_size
+    chip_throughput = system_throughput / world_size
+    # For 70B
+    # all_param_num = getattr(args, "all_param_num", None)
+    # assert all_param_num is not None, "please set all_param_num"
+    # MFU = chip_throughput * 6 * all_param_num * (1 + args.seq_length / (6 * args.hidden_size) ) / 98e12
+    # # tflops_throughput = chip_throughput / float(config.flops_16bit) * 1e12
+    # # logger.info("Throughput(token per chip per second): " + str(chip_throughput))
+    # # logger.info("MFU: " + str(MFU))
+    # # logger.info("Throughput(token per TFLOPS): " + str(tflops_throughput))
+    h = args.hidden_size
+    s = args.seq_length
+    N = 12 * args.num_layers * h **2
+    D = 1
+
+    attn_matmul = 2 * N * D
+    attn_sdp = N * D * (s / h)
+    mlp_matmul = 4 * N * D
+    # moe
+    if args.num_experts is None:
+        factor = 1
+    else:
+        factor = args.moe_router_topk
+    activated_dense_flops = attn_matmul + attn_sdp + mlp_matmul * factor
+    if args.num_experts is not None:
+        act_params = N + args.num_layers *(args.num_experts - 1) * 8 * h**2
+        if torch.distributed.get_rank() == 0:
+            print(f"N: {N} Act param: {act_params} Act flops: {activated_dense_flops}")
+    tflops =  chip_throughput *  activated_dense_flops
+    mfu = tflops / 98e12
+
+    return chip_throughput, mfu
+
+def num_floating_point_operations(args, batch_size):
+    # Attention projection size.
+    query_projection_size = args.kv_channels * args.num_attention_heads
+    query_projection_to_hidden_size_ratio = query_projection_size / args.hidden_size
+    # Group Query Attention.
+    if not args.group_query_attention:
+        args.num_query_groups = args.num_attention_heads
+    # MoE.
+    num_experts_routed_to = 1 if args.num_experts is None else args.moe_router_topk
+    gated_linear_multiplier = 3 / 2 if args.swiglu else 1
+    shared_expert_ffn_hidden_size = (
+        0
+        if args.moe_shared_expert_intermediate_size is None
+        else args.moe_shared_expert_intermediate_size
+    )
+    if not args.multi_latent_attention:
+        return (
+            12
+            * batch_size
+            * args.seq_length
+            * args.num_layers
+            * args.hidden_size
+            * args.hidden_size
+            * (
+                # Attention.
+                (
+                    (
+                        1
+                        + (args.num_query_groups / args.num_attention_heads)
+                        + (args.seq_length / args.hidden_size)
+                    ) * query_projection_to_hidden_size_ratio
+                )
+                # MLP.
+                + (
+                    (args.moe_ffn_hidden_size / args.hidden_size)
+                    * num_experts_routed_to
+                    * gated_linear_multiplier
+                )
+                # Shared Experts.
+                + ((shared_expert_ffn_hidden_size / args.hidden_size) * gated_linear_multiplier)
+                # Logit.
+                + (args.padded_vocab_size / (2 * args.num_layers * args.hidden_size))
+            )
+        )
+    else:
+        if args.q_lora_rank is None:
+            mla_flops_q = args.hidden_size * args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+        else:
+            mla_flops_q = args.hidden_size * args.q_lora_rank +\
+                  args.q_lora_rank * args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+        return (
+            6
+            * batch_size
+            * args.seq_length
+            * args.num_layers
+            * (
+                # MLA Attention.
+                (
+                    (
+                        mla_flops_q
+                        + args.hidden_size * (args.kv_lora_rank + args.qk_pos_emb_head_dim)
+                        + args.num_attention_heads * args.kv_lora_rank * (args.qk_head_dim + args.v_head_dim)
+                        + args.num_attention_heads * args.seq_length * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+                        + args.num_attention_heads * args.seq_length * args.v_head_dim
+                        + args.num_attention_heads * args.v_head_dim * args.hidden_size
+                    )
+                )
+                # Router
+                + args.hidden_size * args.num_experts
+                # MLP.
+                + (
+                    2 * args.hidden_size *  args.moe_ffn_hidden_size * num_experts_routed_to
+                    * gated_linear_multiplier
+                )
+                # Shared Experts.
+                + (2 * args.hidden_size * shared_expert_ffn_hidden_size * gated_linear_multiplier)
+                # Logit.
+                + (args.padded_vocab_size * args.hidden_size / args.num_layers)
+            )
+        )
+
+
+def update_e_score_correction_bias(model, micro_batch_num, ):
+    """
+    Decrease the bias term by gamma if its corresponding expert is overloaded,
+    and increase it by gamma if its corresponding expert is underloaded.
+    """
+    for layer in model.modules():
+        if isinstance(layer, TopKRouter):
+            if layer.moe_noaux_gamma:
+                # allreduce tokens_per_expert
+                layer.tokens_per_expert /= micro_batch_num
+                torch.distributed.all_reduce(layer.tokens_per_expert, op=torch.distributed.ReduceOp.SUM)
+                layer.tokens_per_expert /= torch.distributed.get_world_size()
+                with torch.no_grad():
+                    update_value = torch.where(
+                        layer.tokens_per_expert > 1,
+                        -layer.moe_noaux_gamma,
+                        layer.moe_noaux_gamma
+                    )
+                    layer.e_score_correction_bias += update_value
+                    layer.tokens_per_expert.zero_()
+
+def training_log(loss_dict, total_loss_dict, learning_rate, decoupled_learning_rate, iteration,
+                 loss_scale, report_memory_flag, skipped_iter,
+                 grad_norm, params_norm, num_zeros_in_grad):
+    """Log training information such as losses, timing, ...."""
+    args = get_args()
+    timers = get_timers()
+    writer = get_tensorboard_writer()
+    wandb_writer = get_wandb_writer()
+    one_logger = get_one_logger()
+
+    # Advanced, skipped, and Nan iterations.
+    advanced_iters_key = 'advanced iterations'
+    skipped_iters_key = 'skipped iterations'
+    nan_iters_key = 'nan iterations'
+    # Advanced iterations.
+    if not skipped_iter:
+        total_loss_dict[advanced_iters_key] = total_loss_dict.get(
+            advanced_iters_key, 0) + 1
+    else:
+        if advanced_iters_key not in total_loss_dict:
+            total_loss_dict[advanced_iters_key] = 0
+    # Skipped iterations.
+    total_loss_dict[skipped_iters_key] = total_loss_dict.get(
+        skipped_iters_key, 0) + skipped_iter
+    # Update losses and set nan iterations
+    got_nan = False
+    for key in loss_dict:
+        if not skipped_iter:
+            total_loss_dict[key] = total_loss_dict.get(
+                key, torch.tensor([0.0], dtype=torch.float, device='cuda')) + loss_dict[key]
+        else:
+            value = loss_dict[key].float().sum().item()
+            is_nan = value == float('inf') or \
+                     value == -float('inf') or \
+                     value != value
+            got_nan = got_nan or is_nan
+    total_loss_dict[nan_iters_key] = total_loss_dict.get(
+        nan_iters_key, 0) + int(got_nan)
+
+    # Logging.
+    timers_to_log = [
+        'forward-backward',
+        'forward-compute',
+        'backward-compute',
+        'batch-generator',
+        'forward-recv',
+        'forward-send',
+        'backward-recv',
+        'backward-send',
+        'forward-send-forward-recv',
+        'forward-send-backward-recv',
+        'backward-send-forward-recv',
+        'backward-send-backward-recv',
+        'forward-backward-send-forward-backward-recv',
+        'layernorm-grads-all-reduce',
+        'embedding-grads-all-reduce',
+        'all-grads-sync',
+        'params-all-gather',
+        'optimizer-copy-to-main-grad',
+        'optimizer-unscale-and-check-inf',
+        'optimizer-clip-main-grad',
+        'optimizer-count-zeros',
+        'optimizer-inner-step',
+        'optimizer-copy-main-to-model-params',
+        'optimizer']
+
+    # Calculate batch size.
+    batch_size = args.micro_batch_size * args.data_parallel_size * \
+        get_num_microbatches()
+
+    # Track app tag & app tag ID
+    one_logger_utils.track_app_tag(batch_size, args.world_size, args.seq_length)
+
+    total_iterations = total_loss_dict[advanced_iters_key] + \
+                       total_loss_dict[skipped_iters_key]
+
+    # Tensorboard values.
+    # Timer requires all the ranks to call.
+    if args.log_timers_to_tensorboard and \
+       (iteration % args.tensorboard_log_interval == 0):
+        timers.write(timers_to_log, writer, iteration,
+                     normalizer=total_iterations)
+    if writer and (iteration % args.tensorboard_log_interval == 0):
+        if wandb_writer:
+            wandb_writer.log({'samples vs steps': args.consumed_train_samples},
+                             iteration)
+        writer.add_scalar('learning-rate', learning_rate, iteration)
+        if args.decoupled_lr is not None:
+            writer.add_scalar('decoupled-learning-rate', decoupled_learning_rate, iteration)
+        writer.add_scalar('learning-rate vs samples', learning_rate,
+                          args.consumed_train_samples)
+        if wandb_writer:
+            wandb_writer.log({'learning-rate': learning_rate}, iteration)
+        if args.skipped_train_samples > 0:
+            writer.add_scalar('skipped-train-samples', args.skipped_train_samples, iteration)
+            if wandb_writer:
+                wandb_writer.log({'skipped-train-samples': args.skipped_train_samples}, iteration)
+        writer.add_scalar('batch-size', batch_size, iteration)
+        writer.add_scalar('batch-size vs samples', batch_size,
+                          args.consumed_train_samples)
+        if wandb_writer:
+            wandb_writer.log({'batch-size': batch_size}, iteration)
+        for key in loss_dict:
+            writer.add_scalar(key , loss_dict[key], iteration)
+            writer.add_scalar(key + ' vs samples', loss_dict[key],
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({key: loss_dict[key]}, iteration)
+        if args.log_loss_scale_to_tensorboard:
+            writer.add_scalar('loss-scale', loss_scale, iteration)
+            writer.add_scalar('loss-scale vs samples', loss_scale,
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({'loss-scale': loss_scale}, iteration)
+        if args.log_world_size_to_tensorboard:
+            writer.add_scalar('world-size', args.world_size, iteration)
+            writer.add_scalar('world-size vs samples', args.world_size,
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({'world-size': args.world_size}, iteration)
+        if grad_norm is not None:
+            writer.add_scalar('grad-norm', grad_norm, iteration)
+            writer.add_scalar('grad-norm vs samples', grad_norm,
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({'grad-norm': grad_norm}, iteration)
+        if num_zeros_in_grad is not None:
+            writer.add_scalar('num-zeros', num_zeros_in_grad, iteration)
+            writer.add_scalar('num-zeros vs samples', num_zeros_in_grad,
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({'num-zeros': num_zeros_in_grad}, iteration)
+        if params_norm is not None:
+            writer.add_scalar('params-norm', params_norm, iteration)
+            writer.add_scalar('params-norm vs samples', params_norm,
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({'params-norm': params_norm}, iteration)
+        if args.log_memory_to_tensorboard:
+            mem_stats = torch.cuda.memory_stats()
+            writer.add_scalar(
+                "mem-reserved-bytes",
+                mem_stats["reserved_bytes.all.current"],
+                iteration,
+            )
+            writer.add_scalar(
+                "mem-allocated-bytes",
+                mem_stats["allocated_bytes.all.current"],
+                iteration,
+            )
+            writer.add_scalar(
+                "mem-allocated-count",
+                mem_stats["allocation.all.current"],
+                iteration,
+            )
+    if args.num_experts is not None:
+        moe_loss_scale = 1 / get_num_microbatches()
+        track_moe_metrics(moe_loss_scale, iteration, writer, wandb_writer, total_loss_dict, args.moe_per_layer_logging)
+
+    if iteration % args.log_interval == 0:
+        elapsed_time = timers('interval-time').elapsed(barrier=True)
+        elapsed_time_per_iteration = elapsed_time / total_iterations
+
+        throughput = num_floating_point_operations(args, batch_size) / (
+            elapsed_time_per_iteration * 10**12 * args.world_size)
+
+        one_logger_utils.track_e2e_metrics(args.log_throughput, throughput)
+
+        if args.log_timers_to_tensorboard:
+            if writer:
+                writer.add_scalar('iteration-time',
+                                  elapsed_time_per_iteration, iteration)
+            if wandb_writer:
+                wandb_writer.log({'iteration-time': elapsed_time_per_iteration},
+                                 iteration)
+        log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
+        log_string += ' iteration {:8d}/{:8d} |'.format(
+            iteration, args.train_iters)
+        # chip_throughput, mfu = throughput_calculator(args, elapsed_time_per_iteration, batch_size * args.seq_length)
+        # log_string += ' chip_throughput: {:.2f} /s |'.format(chip_throughput)
+        # log_string += ' mfu: {:.4f} |'.format(mfu)
+        log_string += ' consumed samples: {:12d} |'.format(
+            args.consumed_train_samples)
+        if args.skipped_train_samples > 0:
+            log_string += ' skipped samples: {:12d} |'.format(
+                args.skipped_train_samples)
+        log_string += ' elapsed time per iteration (ms): {:.1f} |'.format(
+            elapsed_time_per_iteration * 1000.0)
+        if args.log_throughput:
+            log_string += f' throughput per GPU (TFLOP/s/GPU): {throughput:.1f} |'
+            if args.log_timers_to_tensorboard:
+                if writer:
+                    writer.add_scalar('throughput', throughput, iteration)
+                if wandb_writer:
+                    wandb_writer.log({'throughput': throughput}, iteration)
+        assert learning_rate is not None
+        # Decoupled_learning_rate should be not None only on first and last pipeline stage.
+        log_string += ' learning rate: {:.6E} |'.format(learning_rate)
+        if args.decoupled_lr is not None and (mpu.is_pipeline_first_stage(ignore_virtual=True) or
+                                              mpu.is_pipeline_last_stage(ignore_virtual=True)):
+            assert decoupled_learning_rate is not None
+            log_string += ' decoupled learning rate: {:.6E} |'.format(decoupled_learning_rate)
+        else:
+            assert decoupled_learning_rate is None
+        log_string += ' global batch size: {:5d} |'.format(batch_size)
+        for key in total_loss_dict:
+            if key not in [advanced_iters_key, skipped_iters_key,
+                           nan_iters_key]:
+                avg = total_loss_dict[key].item() / \
+                      float(max(1, total_loss_dict[advanced_iters_key]))
+                if avg > 0.0:
+                    log_string += ' {}: {:.6E} |'.format(key, avg)
+                total_loss_dict[key] = torch.tensor([0.0], dtype=torch.float, device='cuda')
+        log_string += ' loss scale: {:.1f} |'.format(loss_scale)
+        if grad_norm is not None:
+            log_string += ' grad norm: {:.3f} |'.format(grad_norm)
+        if num_zeros_in_grad is not None:
+            log_string += ' num zeros: {:.1f} |'.format(num_zeros_in_grad)
+        if params_norm is not None:
+            log_string += ' params norm: {:.3f} |'.format(params_norm)
+        log_string += ' number of skipped iterations: {:3d} |'.format(
+            total_loss_dict[skipped_iters_key])
+        log_string += ' number of nan iterations: {:3d} |'.format(
+            total_loss_dict[nan_iters_key])
+        total_loss_dict[advanced_iters_key] = 0
+        total_loss_dict[skipped_iters_key] = 0
+        total_loss_dict[nan_iters_key] = 0
+        print_rank_last(log_string)
+        if report_memory_flag and learning_rate > 0.:
+            # Report memory after optimizer state has been initialized.
+            if torch.distributed.get_rank() == 0:
+                num_microbatches = get_num_microbatches()
+                report_theoretical_memory(args, num_microbatches=num_microbatches, verbose=True)
+            report_memory('(after {} iterations)'.format(iteration))
+            report_memory_flag = False
+        timers.log(timers_to_log, normalizer=args.log_interval)
+
+    return report_memory_flag
+
+def train(forward_step_func, model, optimizer, opt_param_scheduler,
+          train_data_iterator, valid_data_iterator,
+          process_non_loss_data_func, config, checkpointing_context, non_loss_data_func):
+    """Train the model function."""
+    args = get_args()
+    timers = get_timers()
+    one_logger = get_one_logger()
+
+    # Write args to tensorboard
+    write_args_to_tensorboard()
+
+    # Turn on training mode which enables dropout.
+    for model_module in model:
+        model_module.train()
+
+    # Tracking loss.
+    total_loss_dict = {}
+
+    # Iterations.
+    iteration = args.iteration
+
+    # Track E2E metrics at the start of training
+    one_logger_utils.on_train_start(iteration=iteration, consumed_train_samples=args.consumed_train_samples,
+                                    train_samples=args.train_samples, seq_length=args.seq_length,
+                                    train_iters=args.train_iters, save=args.save, async_save=args.async_save,
+                                    log_throughput=args.log_throughput,
+                                    num_floating_point_operations_so_far=args.num_floating_point_operations_so_far)
+
+    num_floating_point_operations_so_far = args.num_floating_point_operations_so_far
+
+    # Setup some training config params
+    config.grad_scale_func = optimizer.scale_loss
+    config.timers = timers
+    if isinstance(model[0], DDP) and args.overlap_grad_reduce:
+        assert config.no_sync_func is None, \
+            ('When overlap_grad_reduce is True, config.no_sync_func must be None; '
+             'a custom no_sync_func is not supported when overlapping grad-reduce')
+        config.no_sync_func = [model_chunk.no_sync for model_chunk in model]
+        if len(model) == 1:
+            config.no_sync_func = config.no_sync_func[0]
+        if args.align_grad_reduce:
+            config.grad_sync_func = [model_chunk.start_grad_sync for model_chunk in model]
+            if len(model) == 1:
+                config.grad_sync_func = config.grad_sync_func[0]
+    if args.overlap_param_gather and args.align_param_gather:
+        config.param_sync_func = [model_chunk.start_param_sync for model_chunk in model]
+        if len(model) == 1:
+            config.param_sync_func = config.param_sync_func[0]
+    config.finalize_model_grads_func = finalize_model_grads
+
+    timers('interval-time', log_level=0).start(barrier=True)
+    print_datetime('before the start of training step')
+    report_memory_flag = True
+    # exit = False
+    pre_hook_enabled = False
+    should_exit = False
+    exit_code = 0
+
+    if args.manual_gc:
+        # Disable the default garbage collector and perform the collection manually.
+        # This is to align the timing of garbage collection across ranks.
+        assert args.manual_gc_interval >= 0, \
+            'Manual garbage collection interval should be laerger than or equal to 0.'
+        gc.disable()
+        gc.collect()
+
+    # Singleton Initialization
+    if args.log_straggler:
+        global stimer
+        world = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+        mmcnt = args.straggler_minmax_count
+        stimer.configure(world, rank,
+                mmcnt = mmcnt,
+                enabled = not args.disable_straggler_on_startup,
+                port = args.straggler_ctrlr_port)
+    # total_flops = 0.0
+    num_floating_point_operations_since_last_log_event = 0.0
+
+    num_microbatches = get_num_microbatches()
+    eval_duration = 0.0
+    eval_iterations = 0
+
+    def get_e2e_base_metrics():
+        """Get base metrics values for one-logger to calculate E2E tracking metrics.
+        """
+        return {
+            'iteration': iteration,
+            'train_duration': timers('interval-time').active_time(),
+            'eval_duration': eval_duration,
+            'eval_iterations': eval_iterations,
+            'total_flops': num_floating_point_operations_since_last_log_event,
+            'num_floating_point_operations_so_far': num_floating_point_operations_so_far,
+            'consumed_train_samples': args.consumed_train_samples,
+            'world_size': args.world_size,
+            'seq_length': args.seq_length
+        }
+    # Cache into one-logger for callback
+    if one_logger:
+        with one_logger.get_context_manager():
+            one_logger.store_set('get_e2e_base_metrics', get_e2e_base_metrics)
+
+    prof = None
+    if args.profile and torch.distributed.get_rank() in args.profile_ranks and args.use_pytorch_profiler:
+        prof = torch.profiler.profile(
+        schedule=torch.profiler.schedule(
+            wait=max(args.profile_step_start-1, 0),
+            warmup=1 if args.profile_step_start > 0 else 0,
+            active=args.profile_step_end-args.profile_step_start,
+            repeat=1),
+        on_trace_ready=torch.profiler.tensorboard_trace_handler(args.tensorboard_dir),
+        record_shapes=True,
+        with_stack=True)
+        prof.start()
+
+    start_iteration = iteration
+    # Disable forward pre-hook to start training to ensure that errors in checkpoint loading
+    # or random initialization don't propagate to all ranks in first all-gather (which is a
+    # no-op if things work correctly).
+    if args.use_distributed_optimizer and args.overlap_param_gather:
+        disable_forward_pre_hook(model, param_sync=False)
+        # Also remove param_sync_func temporarily so that sync calls made in
+        # `forward_backward_func` are no-ops.
+        param_sync_func = config.param_sync_func
+        config.param_sync_func = None
+        pre_hook_enabled = False
+    # Also, check weight hash across DP replicas to be very pedantic.
+    if args.check_weight_hash_across_dp_replicas_interval is not None:
+        assert check_param_hashes_across_dp_replicas(model, cross_check=True), \
+            "Parameter hashes not matching across DP replicas"
+        torch.distributed.barrier()
+        print_rank_0(f">>> Weight hashes match after {iteration} iterations...")
+
+    with maybe_enable_profiling(
+        args, global_step=iteration
+    ) as torch_profiler:
+        while iteration < args.train_iters:
+            if args.profile and torch.distributed.get_rank() in args.profile_ranks:
+                if args.use_pytorch_profiler:
+                    prof.step()
+                elif iteration == args.profile_step_start:
+                    torch.cuda.cudart().cudaProfilerStart()
+                    torch.autograd.profiler.emit_nvtx(record_shapes=True).__enter__()
+
+            maybe_finalize_async_save(blocking=False)
+
+            # Update number of microbatches first without consistency check to decide if a
+            # checkpoint should be saved. If the number of microbatches is different
+            # from the previous iteration, save a checkpoint. Then run consistency check
+            # to make sure training configuration is still valid.
+            update_num_microbatches(args.consumed_train_samples, consistency_check=False, verbose=True)
+            if get_num_microbatches() != num_microbatches and iteration != 0:
+                assert get_num_microbatches() > num_microbatches, \
+                    "number of microbatches should be increasing due to batch size rampup ... %d -> %d." % (num_microbatches, get_num_microbatches())
+                if args.save is not None:
+                    save_checkpoint_and_time(iteration, model, optimizer,
+                                            opt_param_scheduler,
+                                            num_floating_point_operations_so_far,
+                                            checkpointing_context, train_data_iterator=train_data_iterator)
+            num_microbatches = get_num_microbatches()
+            update_num_microbatches(args.consumed_train_samples, consistency_check=True, verbose=True)
+
+            args.curr_iteration = iteration
+            loss_dict, skipped_iter, should_checkpoint, should_exit, exit_code, grad_norm, num_zeros_in_grad = \
+                train_step(forward_step_func,
+                        train_data_iterator,
+                        model,
+                        optimizer,
+                        opt_param_scheduler,
+                        config)
+            if config.moe_noaux_gamma:
+                update_e_score_correction_bias(model[0], num_microbatches)
+
+            if should_checkpoint:
+                save_checkpoint_and_time(iteration, model, optimizer,
+                                        opt_param_scheduler,
+                                        num_floating_point_operations_so_far,
+                                        checkpointing_context, train_data_iterator=train_data_iterator)
+            if should_exit:
+                break
+
+            # Enable forward pre-hooks after first set of forward and backward passes.
+            # When running in fp16, skip all NaN iterations until steady-state loss scaling value
+            # is reached.
+            if iteration == start_iteration:
+                if skipped_iter:
+                    # Only enable forward pre-hook after a training step has successfully run. Relevant
+                    # for fp16 codepath where first XX iterations are skipped until steady-state loss
+                    # scale value is reached.
+                    start_iteration = iteration + 1
+                else:
+                    # Enable forward pre-hook after training step has successfully run. All subsequent
+                    # forward passes will use the forward pre-hook / `param_sync_func` in
+                    # `forward_backward_func`.
+                    if args.use_distributed_optimizer and args.overlap_param_gather:
+                        enable_forward_pre_hook(model)
+                        config.param_sync_func = param_sync_func
+                        pre_hook_enabled = True
+
+            if torch_profiler:
+                torch_profiler.step()
+            iteration += 1
+            batch_size = mpu.get_data_parallel_world_size() * \
+                        args.micro_batch_size * \
+                        get_num_microbatches()
+            args.consumed_train_samples += batch_size
+            num_skipped_samples_in_batch = (get_current_global_batch_size() -
+                                            get_current_running_global_batch_size())
+            if args.decrease_batch_size_if_needed:
+                assert num_skipped_samples_in_batch >= 0
+            else:
+                assert num_skipped_samples_in_batch == 0
+            args.skipped_train_samples += num_skipped_samples_in_batch
+            num_floating_point_operations_in_batch = num_floating_point_operations(args, batch_size)
+            num_floating_point_operations_so_far += num_floating_point_operations_in_batch
+            num_floating_point_operations_since_last_log_event += num_floating_point_operations_in_batch
+
+            # # Send heartbeat to FT package and update timeouts.
+            # if args.enable_ft_package:
+            #     ft_client = ft_integration.get_rank_monitor_client(
+            #         ft_integration.StateMachineActions.TRAIN_HEARTBEAT)
+            #     if ft_client is not None:
+            #         ft_client.send_heartbeat()
+            #         # TODO we are always calculating timeouts in the current implementation
+            #         # if we want to rely on manually setup then we need to add additional argument
+            #         # to training and pass it here
+            #         if ft_integration.can_update_timeouts():
+            #             ft_integration.get_rank_monitor_client(
+            #                 ft_integration.StateMachineActions.UPDATE_TIMEOUT).calculate_and_set_timeouts()
+            #             print_rank_0(f'Updated FT timeouts. New values: \
+            #                 {ft_integration.get_rank_monitor_client().timeouts}')
+
+            # # Bring CPU and GPU back in sync if on right iteration.
+            # if (
+            #     args.train_sync_interval
+            #     and iteration % args.train_sync_interval == 0
+            # ):
+            #     torch.cuda.synchronize()
+
+            # Logging.
+            if not optimizer.is_stub_optimizer:
+                loss_scale = optimizer.get_loss_scale().item()
+            else:
+                loss_scale = 1.0
+            params_norm = None
+            if args.log_params_norm:
+                params_norm = calc_params_l2_norm(model)
+
+            learning_rate = None
+            decoupled_learning_rate = None
+            for param_group in optimizer.param_groups:
+                if param_group['is_decoupled_lr']:
+                    decoupled_learning_rate = param_group['lr']
+                else:
+                    learning_rate = param_group['lr']
+            report_memory_flag = training_log(loss_dict, total_loss_dict,
+                                            learning_rate,
+                                            decoupled_learning_rate,
+                                            iteration, loss_scale,
+                                            report_memory_flag, skipped_iter,
+                                            grad_norm, params_norm, num_zeros_in_grad)
+
+            # # StragglerDetector
+            # if iteration % args.log_interval == 0 and args.log_straggler:
+            #     stimer.report(total_flops, args.log_interval)
+            #     total_flops = 0.0
+
+            # if args.check_weight_hash_across_dp_replicas_interval is not None and \
+            #         iteration % args.check_weight_hash_across_dp_replicas_interval == 0:
+            #     if args.use_distributed_optimizer and args.overlap_param_gather:
+            #         optimizer.disable_pre_hook()
+            #     assert check_param_hashes_across_dp_replicas(model, cross_check=True), \
+            #         "Parameter hashes not matching across DP replicas"
+            #     torch.distributed.barrier()
+            #     print_rank_0(f">>> Weight hashes match after {iteration} iterations...")
+            #     if args.use_distributed_optimizer and args.overlap_param_gather:
+            #         optimizer.enable_pre_hook()
+
+            # # Autoresume
+            # if args.adlr_autoresume and \
+            # (iteration % args.adlr_autoresume_interval == 0):
+            #     check_adlr_autoresume_termination(iteration, model, optimizer,
+            #                                     opt_param_scheduler)
+
+            # Evaluation
+            if args.eval_interval and iteration % args.eval_interval == 0 and \
+                args.do_valid:
+                timers('interval-time').stop()
+                if args.use_distributed_optimizer and args.overlap_param_gather:
+                    disable_forward_pre_hook(model)
+                    pre_hook_enabled = False
+                if args.manual_gc and args.manual_gc_eval:
+                    # Collect all objects.
+                    gc.collect()
+                prefix = f'iteration {iteration}'
+                timers('eval-time', log_level=0).start(barrier=True)
+                evaluate_and_print_results(prefix, forward_step_func,
+                                        valid_data_iterator, model,
+                                        iteration, process_non_loss_data_func,
+                                        config, verbose=False, write_to_tensorboard=True,
+                                        non_loss_data_func=non_loss_data_func)
+                eval_duration += timers('eval-time').elapsed()
+                eval_iterations += args.eval_iters
+                timers('eval-time').stop()
+                one_logger_utils.track_e2e_metrics()
+
+                if args.manual_gc and args.manual_gc_eval:
+                    # Collect only the objects created and used in evaluation.
+                    gc.collect(generation=0)
+                if args.use_distributed_optimizer and args.overlap_param_gather:
+                    enable_forward_pre_hook(model)
+                    pre_hook_enabled = True
+                timers('interval-time', log_level=0).start(barrier=True)
+
+
+                if args.enable_ft_package and ft_integration.get_rank_monitor_client() is not None:
+                    ft_integration.get_rank_monitor_client(
+                        ft_integration.StateMachineActions.EVAL_HEARTBEAT).send_heartbeat()
+
+            # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC).
+            # Some of these only happen at specific iterations.
+            post_training_step_callbacks(model, optimizer, opt_param_scheduler, iteration, prof,
+                                        num_floating_point_operations_since_last_log_event)
+
+            # Checkpoint and decide whether to exit.
+            should_exit = checkpoint_and_decide_exit(model, optimizer, opt_param_scheduler, iteration,
+                                                    num_floating_point_operations_so_far,
+                                                    checkpointing_context, train_data_iterator)
+            if should_exit:
+                break
+
+    one_logger_utils.track_e2e_metrics()
+
+    # Flush TensorBoard, WandB writers and one-logger.
+    writer = get_tensorboard_writer()
+    if writer:
+        writer.flush()
+
+    # Close out pre-hooks if using distributed optimizer and overlapped param gather.
+    if pre_hook_enabled:
+        disable_forward_pre_hook(model)
+
+    if args.enable_ft_package and ft_integration.get_rank_monitor_client() is not None:
+        ft_integration.get_rank_monitor_client().shutdown_workload_monitoring()
+
+    maybe_finalize_async_save(blocking=True)
+
+    # If any exit conditions (signal handler, duration, iterations) have been reached, exit.
+    if should_exit:
+        wandb_writer = get_wandb_writer()
+        if wandb_writer:
+            wandb_writer.finish()
+        sys.exit(exit_code)
+
+    return iteration, num_floating_point_operations_so_far
+
+import megatron.training
+megatron.training.training.training_log = training_log
+
+enable_profiler = int(os.getenv("ENABLE_PROFILER", 0))
+# if enable_profiler:
+megatron.training.training.train = train
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/cuda_patch/transformer_config.py b/megatron-lm-musa-patch/cuda_patch/transformer_config.py
new file mode 100644
index 000000000..16ff872d5
--- /dev/null
+++ b/megatron-lm-musa-patch/cuda_patch/transformer_config.py
@@ -0,0 +1,718 @@
+
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch.nn.functional as F
+
+from megatron.core.transformer.enums import AttnBackend
+
+from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.utils import get_te_version, init_method_normal, is_te_min_version, scaled_init_method_normal
+
+
+@dataclass
+class TransformerConfig(ModelParallelConfig):
+    """Configuration object for megatron-core transformers.
+
+    The initialization function has an argument for each parameter,
+    including those in ModelParallelConfig.
+    """
+
+    ####################
+    # model architecture
+    ####################
+    num_layers: int = 0
+    """Number of transformer layers in a transformer block."""
+
+    first_pipeline_num_layers: int = None
+    """Number of transformer layers on first pipeline stage.
+    None implies equal layer division across PP ranks."""
+
+    last_pipeline_num_layers: int = None
+    """Number of transformer layers on last pipeline stage.
+    None implies equal layer division across PP ranks."""
+
+    hidden_size: int = 0
+    """Transformer hidden size."""
+
+    num_attention_heads: int = 0
+    """Number of transformer attention heads."""
+
+    attention_backend: AttnBackend = AttnBackend.auto
+    """Attention backend to run. By default we let transformer engine
+    decide the best backend to run (except in the case of local).
+    If attention backend is local we use the local pytorch implementation in mcore.
+    Users can specify exact backend by changing this config. """
+
+    softmax_scale: float = None
+    """Softmax scale for attention scaling."""
+
+    num_query_groups: int = None
+    """Number of query groups for group query attention. If None, normal attention is used."""
+
+    ffn_hidden_size: int = None
+    """Transformer Feed-Forward Network hidden size. This is set to 4*hidden_size
+    if not provided."""
+
+    kv_channels: int = None
+    """Projection weights dimension in multi-head attention. This is set to hidden_size //
+    num_attention_heads if not provided."""
+
+    hidden_dropout: float = 0.1
+    """Dropout probability for transformer hidden state."""
+
+    attention_dropout: float = 0.1
+    """Post attention dropout probability."""
+
+    fp32_residual_connection: bool = False
+    """If true, move residual connections to fp32."""
+
+    # @jcasper should we keep this option?
+    apply_residual_connection_post_layernorm: bool = False
+    """If True, uses the original BERT residule connection ordering."""
+
+    layernorm_epsilon: float = 1e-5
+    """Epsilon value for any LayerNorm operations."""
+
+    layernorm_zero_centered_gamma: bool = False
+    """If set to True, the LayerNorm is adjusted to center the gamma values around 0. This improves
+    numerical stability."""
+
+    add_bias_linear: bool = True
+    """Include a bias term in all linear layers (QKV projections, after core attention, and two in
+    MLP layer)."""
+
+    add_qkv_bias: bool = False
+    """Add a bias term only for QKV projections."""
+
+    gated_linear_unit: bool = False
+    """Use a gated linear unit for the first linear layer in the MLP."""
+
+    activation_func: Callable = F.gelu
+    """Activation function to use for the non-linearity in the MLP."""
+
+    activation_func_fp8_input_store: bool = False
+    """Store the input of MLP activation function in FP8 for backprop to save memory.
+    The stored input is casted back to the original precision before backprop compuatation."""
+
+    num_moe_experts: int = None
+    """Number of experts to use for MoE layer. When set, it replaces MLP with MoE layer. Set to None
+    for no MoE."""
+
+    rotary_interleaved: bool = False
+    """True is rotate pairs of even and odd dimensions (RoFormer style), False is rotate pairs of
+    first half and second half (LLaMa style). Default to False."""
+
+    window_size: Optional[Tuple[int, int]] = None
+    """If not None, then will use sliding window attention. The size of the window is specified by
+    the numbers inside the tuple; -1 is special value meaning "infinite window size"."""
+
+    normalization: bool = "LayerNorm"
+    """Which norm to use for normalization layers, valid options are `LayerNorm` and `RMSNorm`."""
+
+    qk_layernorm: bool = False
+    """Whether to apply LayerNorm to the query and key embeddings."""
+
+    test_mode: bool = False
+    """Whether to run real-time tests."""
+
+    calculate_per_token_loss: bool = False
+    """Whether cross entropy loss is calculated over the actual number of non-padded tokens in the
+    global batch, versus the default behavior of assuming all tokens are non-padded."""
+
+    multi_latent_attention: bool = False
+    """Whether to use multi-latent attention."""
+
+    ####################
+    # initialization
+    ####################
+    init_method: Callable = None
+    """Method to initialize weights. Note that bias is always set to zero. Should be a function that
+    takes a single Tensor and initializes it. If None, will be set to
+    megatron.core.utils.init_method_normal(init_method_std) which is torch nn init normal with
+    mean=0.0 and std=init_method_std."""
+
+    output_layer_init_method: Callable = None
+    """Method to initialize weights of the output layer of both attention and MLP blocks. If None,
+    will be set to megatron.core.utils.scaled_init_method_normal(init_method_std) which is torch nn
+    init normal with mean=0.0 and std=init_method_std / math.sqrt(2.0 * num_layers)."""
+
+    init_method_std: float = 0.02
+    """Standard deviation of the zero mean normal for the default initialization method, not used if
+    init_method and output_layer_init_method are provided."""
+
+    ####################
+    # mixed-precision
+    ####################
+    apply_query_key_layer_scaling: bool = False
+    """If true, scale Q * K^T by 1 / layer-number. This improve numeric stability when training with
+    fp16."""
+
+    attention_softmax_in_fp32: bool = True
+    """If True, run attention masking and softmax in fp32. This should be True if
+    apply_query_key_layer_scaling is True."""
+
+    ####################
+    # fusion
+    ####################
+    bias_activation_fusion: bool = False
+    """If True, fuses bias addition and the activation function when possible."""
+
+    masked_softmax_fusion: bool = False
+    """If True, uses softmax fusion."""
+
+    persist_layer_norm: bool = False
+    """If True, uses the persistent fused layer norm kernel. This kernel only supports a fixed set
+    of hidden sizes."""
+
+    memory_efficient_layer_norm: bool = False
+    """If True, and using local layers (not from TransformerEngine), tells Apex to use the memory
+    efficient fused LayerNorm kernel. Ignored if not using LayerNorm."""
+
+    bias_dropout_fusion: bool = False  # TODO: this should be bias_dropout_add_fusion?
+    """If True, uses bias dropout fusion."""
+
+    apply_rope_fusion: bool = False
+    """If True, use fused RoPE kernel."""
+
+    ####################
+    # activation recomputation
+    ####################
+    recompute_granularity: str = None
+    """Determines which type of activation recompute to use.  Megatron-core supports 'selective'
+    activation checkpointing where only the memory intensive part of attention is checkpointed.
+    These memory intensive activations are also less compute intensive which makes activation
+    checkpointing more efficient for LLMs (20B+).  See Reducing Activation Recomputation in Large
+    Transformer Models (https://arxiv.org/abs/2205.05198) for more details.  'full' will checkpoint
+    the entire transformer layer.  If None, no recompute is performed and all activations are saved.
+    If set, must be 'selective' or 'full'. 'selective' always uses all layers.
+    """
+
+    recompute_method: str = None
+    """Determines which transformer layers will be recomputed. uniform will uniformly divide the
+    total number of transformer layers in a transformer block and recompute the input activation of
+    each divided chunk at the specified granularity.  block will recompute the input activations for
+    only a set number of transformer layers per pipeline stage.  The rest of the layers in the
+    pipeline stage will not have any activations recomputed.  If None, and recompute is enabled, all
+    layers will do recomputation. If set, must be 'uniform' or 'block'."""
+
+    recompute_num_layers: int = None
+    """When recompute_method is uniform, recompute_num_layers is the number of transformer layers in
+    each uniformly divided recompute unit.  When recompute_method is block, recompute_num_layers is
+    the number of transformer layers to recompute within each pipeline stage.  Must be None for
+    'selective' activation checkpointing."""
+
+    distribute_saved_activations: bool = None
+    """If True, distribute recomputed activations across the model parallel group."""
+
+    ####################
+    # fp8 related
+    ####################
+    fp8: str = None
+    """If set, enables the use of FP8 precision through Transformer Engine. There are 2 predefined
+    choices (1) 'e4m3' uniformly uses e4m3 for all FP8 tensors, (2) 'hybrid' uses e4m3 for all FP8
+    activation and weight tensors and e5m2 for all FP8 output activation gradient tensors."""
+
+    fp8_margin: int = 0
+    """Margin for the scaling factor computation."""
+
+    fp8_interval: int = 1
+    """DEPRECATED from TransformerEngine v1.8.0. This flag is ignored.
+    Controls how often the scaling factor is recomputed.
+    """
+
+    fp8_amax_history_len: int = 1
+    """The length of the amax history window used for scaling factor computation."""
+
+    fp8_amax_compute_algo: str = "most_recent"
+    """Algorithm used for choosing the `amax` value for the scaling factor computation. There are 2
+    predefined choices: `max` chooses the largest `amax` in the history window, while `most_recent`
+    always chooses the most recently seen value.
+
+    """
+
+    fp8_wgrad: bool = True
+    """When set to False, override FP8 config options and do the wgrad computation
+    in higher precision."""
+
+    fp8_dot_product_attention: bool = False
+    """When set to True, use the FP8 implementation of Dot Product Attention."""
+
+    fp8_multi_head_attention: bool = False
+    """When set to True, use the FP8 implementation of Multi Head Attention."""
+
+    tp_only_amax_red: bool = False
+    """When set to True, reduce the FP8 AMAX only in the TP or TP-CP domain"""
+
+    ####################
+    # MoE related
+    ####################
+    moe_shared_expert_intermediate_size: int = None
+    """Shared expert total ffn hidden size.
+    It should be equal to 'num_shared_experts * ffn_size_of_each_shared_expert' if
+    there are multiple shared experts.
+    None means no shared expert."""
+
+    moe_shared_expert_overlap: bool = False
+    """Enable overlapping between shared expert computations and dispatcher communications.
+    Without this, the shared epxerts execute after the routed experts."""
+
+    moe_layer_freq: int = 1
+    """Frequency between MoE layers and Dense layers. Accepts either:
+    - An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers.
+    - A string containing a Python list expression that defines a custom pattern, e.g.:
+    "([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0]
+    where 1 indicates an expert layer and 0 indicates a dense layer."""
+
+    moe_ffn_hidden_size: int = None
+    """MoE Feed-Forward Network hidden size"""
+
+    moe_router_load_balancing_type: str = "aux_loss"
+    """The load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss
+    used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the loss used in DeepSeekV2,
+    which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing
+    algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss"."""
+
+    moe_router_topk: int = 2
+    """Number of experts to route to for each token."""
+
+    moe_router_topk_limited_devices: int = None
+    """Number of expert parallel ranks to consider for each token during routing. Perform top-k
+    routing on a subset of expert parallel ranks by first selecting N ranks for each token, then
+    conducting top-k selection among experts on these devices. None means no device limitation."""
+
+    moe_router_num_node_group: int = None
+    """Number of node groups for MoE. If None, the number of node groups is equal to the number of
+    expert model parallel groups/8."""
+
+    moe_noaux_gamma: float = None
+    """The gamma value for the noaux loss. Default is None, which means no noaux loss. Recommended set to 1e-3."""
+
+    moe_router_pre_softmax: bool = False
+    """Enable pre-softmax routing for MoE, which means softmax is before the top-k selection.
+    By default, softmax is done after top-k."""
+
+    moe_router_use_sigmoid: bool = False
+    """Use sigmoid instead of softmax for routing scores."""
+
+    moe_router_norm_topk_prob: bool = False
+    """Normalize the top-k probabilities to sum to 1."""
+
+    moe_complementary_seq_aux_loss: bool = False
+    """Use complementary sequence aux loss for MoE."""
+
+    moe_router_topk_scaling_factor: float = None
+    """Scaling factor for routing score in top-k selection, only works when moe_router_pre_softmax
+    enabled. Defaults to None, which means no scaling."""
+
+    moe_grouped_gemm: bool = False
+    """When there are multiple experts per rank, compress multiple local (potentially small) gemms
+    in a single kernel launch to improve the utilization and performance by leveraging the Grouped
+    GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm).
+    """
+
+    moe_use_legacy_grouped_gemm: bool = False
+    """Use legacy GroupedMLP rather than TEGroupedMLP.
+    Note: The legacy one will be deprecated soon."""
+
+    moe_aux_loss_coeff: float = 0  # 1e-2 would be a good start value for load balance loss.
+    """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended."""
+
+    moe_device_level_aux_loss_coeff: float = None
+    """Scaling coefficient for the device_level aux loss. """
+
+    moe_comm_aux_loss_coeff: float = None
+    """Scaling coefficient for the comm aux loss. """
+
+    moe_z_loss_coeff: float = None  # 1e-3 would be a good start value for z-loss
+    """Scaling coefficient for the z-loss. A starting value of 1e-3 is recommended."""
+
+    moe_input_jitter_eps: float = None
+    """Add noise to the input tensor by applying jitter with a specified epsilon value."""
+
+    moe_token_dropping: bool = False  # TODO: Support token dropping.
+    """This feature involves selectively dropping and padding tokens for each expert to achieve a
+    specified capacity, similar to GShard, Switch-Transformer, and DeepSpeed-MoE. Note that this is
+    currently unsupported so should remain False."""
+
+    moe_token_dispatcher_type: str = "allgather"
+    """The type of token dispatcher to use. The default is 'allgather'.
+    Options are 'allgather' and 'alltoall'."""
+
+    moe_per_layer_logging: bool = False
+    """Enable per-layer logging for MoE, currently supports auxiliary loss and z loss."""
+
+    moe_expert_capacity_factor: float = None
+    """moe_expert_capacity_factor (float): The capacity factor for each expert, None means no token
+    will be dropped. The default is None."""
+
+    moe_device_level_capacity: bool = False
+    """moe_device_level_capacity (bool): Whether to consider the expert capacity
+    of a group together. The default is False."""
+
+    moe_pad_expert_input_to_capacity: bool = False
+    """moe_pad_expert_input_to_capacity (bool): If True, pads the input for each expert to match
+    the expert capacity length, effective only after the moe_expert_capacity_factor is set. The
+    default setting is False."""
+
+    moe_token_drop_policy: str = 'probs'
+    """The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with
+    the lowest probabilities will be dropped. If "position", tokens at the end of each batch will
+    be dropped.
+    """
+
+    moe_layer_recompute: bool = False
+    """Memory optimization: checkpointing moe_layer to save actiavtion memory."""
+
+    ##################
+    # Context Parallel
+    ##################
+    cp_comm_type: Union[str, List[str]] = None
+    """Inter-gpu communication type for context parallelism.
+    str: all layers share same communication type.
+    List[str]: each layer has its separate communication type.
+    cp_comm_type of each layer can be "p2p" or "all_gather" or "a2a" or "a2a+p2p".
+    "p2p": Exchange KV chunks with P2P communications in ring topology. P2P is async and can be
+    overlapped with attention compute.
+    "all_gather": All-gather to get full sequence of KV before attention. The all-gather is not
+    async, and cannot be overlapped.
+    "a2a": Like DeepSpeed Ulysses, scatter attention heads across the CP group, and gather to get
+    full sequence of QKV.
+    "a2a+p2p": A hierarchical implementation of context parallelism to attention.
+    It uses A2A communications in low-level CP groups (e.g., via NVLink),
+    and P2P communications in high-level CP groups (e.g., via IBLink).
+    """
+
+    ####################
+    # miscellaneous
+    ####################
+    clone_scatter_output_in_embedding: bool = True
+    """When set to True, clone the output of scatter_to_sequence_parallel_region in embedding layer
+    to facilitate garbage collection of input."""
+
+    disable_parameter_transpose_cache: bool = False
+    """When set to true, the parameter transposes are not cached for subsequent iterations."""
+
+    enable_cuda_graph: bool = False
+    """When set to true, TransformerLayer layers are swapped with a CUDA graphed version."""
+
+    cuda_graph_retain_backward_graph: bool = False
+    """When set to true, cudagraph backward passes will be graph captured with 'retain_grad=True'
+    This may enable cudagraphs for certain modules that are not completely cudagraph safe. For
+    more details, see: https://pytorch.org/docs/stable/generated/torch.Tensor.backward.html."""
+
+    external_cuda_graph: bool = False
+    """When set to true, TransformerLayer layers are swapped with user provided CUDA graphs."""
+
+    config_logger_dir: str = ""
+    """When non-empty, dumps entry-point configs to config_logger_dir"""
+
+    flash_decode: bool = False
+    """ Use the optimized flash decoding kernel during inference. """
+
+    use_te_rng_tracker: bool = False
+    """ Whether to use the TE or MCore version of the RNG tracker. """
+
+    inference_rng_tracker: bool = False
+    """ Whether we should instantiate a separate RNG tracker for inference. """
+
+    def __post_init__(self):
+        """Python dataclass method that is used to modify attributes after initialization.
+        See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more
+        details.
+        """
+        super().__post_init__()
+        if self.fp16 and self.bf16:
+            raise ValueError(
+                f'Only one of self.fp16: {self.fp16} and self.bf16 {self.bf16} should be True.'
+            )
+
+        if self.num_attention_heads % self.tensor_model_parallel_size != 0:
+            raise ValueError(
+                f"num_attention_heads ({self.num_attention_heads}) must be a multiple of "
+                f"tensor_model_parallel_size ({self.tensor_model_parallel_size})."
+            )
+
+        if self.ffn_hidden_size is None:
+            self.ffn_hidden_size = 4 * self.hidden_size
+
+        if self.kv_channels is None:
+            self.kv_channels = self.hidden_size // self.num_attention_heads
+
+        if self.num_query_groups is None:
+            self.num_query_groups = self.num_attention_heads
+
+        if self.num_query_groups % self.tensor_model_parallel_size != 0:
+            raise ValueError(
+                f"num_query_groups ({self.num_query_groups}) must be a multiple of "
+                f"tensor_model_parallel_size ({self.tensor_model_parallel_size})."
+            )
+
+        if self.apply_query_key_layer_scaling:
+            self.attention_softmax_in_fp32 = True
+
+        if self.expert_model_parallel_size > 1 and self.num_moe_experts is None:
+            raise ValueError('num_moe_experts must be non None to use expert-parallel.')
+
+        if self.num_moe_experts is not None and self.num_moe_experts <= 0:
+            raise ValueError('num_moe_experts must be non-negative.')
+
+        if self.moe_ffn_hidden_size is None:
+            self.moe_ffn_hidden_size = self.ffn_hidden_size
+
+        if self.moe_shared_expert_intermediate_size is not None:
+            if self.moe_shared_expert_intermediate_size <= 0:
+                raise ValueError(
+                    f'moe_shared_expert_intermediate_size must be '
+                    f'num_shared_experts * ffn_size_of_each_shared_expert, '
+                    f'but got {self.moe_shared_expert_intermediate_size}'
+                )
+            if self.moe_shared_expert_overlap and self.moe_token_dispatcher_type not in [
+                "alltoall"
+            ]:
+                raise ValueError(
+                    f'moe_shared_expert_overlap only works with alltoall token dispatcher.'
+                )
+
+        if self.moe_expert_capacity_factor is not None:
+            if self.moe_token_dispatcher_type not in ["alltoall", "alltoall_seq"]:
+                raise ValueError(
+                    'moe_expert_capacity_factor only works with alltoall token dispatcher'
+                )
+            if self.moe_expert_capacity_factor < 0:
+                self.moe_expert_capacity_factor = None
+            if self.moe_router_load_balancing_type not in ["aux_loss", "seq_aux_loss", "none"]:
+                raise ValueError(
+                    'moe_expert_capacity_factor only works with aux_loss or none load balancing'
+                )
+
+        if self.moe_pad_expert_input_to_capacity:
+            if self.moe_expert_capacity_factor is None:
+                raise ValueError(
+                    'moe_expert_capacity_factor must be set to use moe_pad_expert_input_to_capacity'
+                )
+
+        if self.cpu_offloading and (
+            self.cpu_offloading_num_layers < 0 or self.cpu_offloading_num_layers >= self.num_layers
+        ):
+            raise ValueError(
+                f'CPU offloading can be done only for layers less than {self.num_layers}'
+            )
+
+        if self.cpu_offloading and self.pipeline_model_parallel_size > 1:
+            raise ValueError(
+                'Currently there is no support for Pipeline parallelism with CPU offloading'
+            )
+
+        if self.cpu_offloading and self.recompute_granularity is not None:
+            raise ValueError(
+                'CPU offloading does not work when activation recomputation is enabled'
+            )
+
+        if self.recompute_granularity is not None:
+            if self.recompute_granularity not in ['full', 'selective']:
+                raise ValueError(
+                    f'When using recompute_granuarlity: {self.recompute_granularity} must be "full"'
+                    'or "selective".'
+                )
+
+            if self.recompute_method is not None:
+                if self.recompute_method not in ['block', 'uniform']:
+                    raise ValueError(
+                        f'recompute_method: {self.recompute_method} must be "block" or "uniform".'
+                    )
+            elif self.recompute_granularity != 'selective':
+                raise ValueError(
+                    f'Using recompute_granularity: {self.recompute_granularity} so '
+                    'recompute_method must be "block" or "uniform"'
+                )
+
+            if self.recompute_granularity != 'selective' and self.recompute_num_layers is None:
+                raise ValueError(
+                    f'When using recompute_granularity: {self.recompute_granularity} '
+                    'recompute_num_layers must be between '
+                    '1 and num_layers_per_pipeline_rank: '
+                    f'{self.num_layers // self.pipeline_model_parallel_size}'
+                )
+            elif (
+                self.recompute_granularity == 'selective' and self.recompute_num_layers is not None
+            ):
+                raise ValueError(
+                    f'When using recompute_granularity: {self.recompute_granularity} '
+                    'recompute_num_layers must be None.'
+                )
+
+            if self.distribute_saved_activations and self.sequence_parallel:
+                raise ValueError(
+                    f'distribute_saved_activations: {self.distribute_saved_activations} must be '
+                    f'false when sequence parallel is enabled: {self.sequence_parallel}'
+                )
+
+            if self.virtual_pipeline_model_parallel_size is not None:
+                if not self.num_layers % self.virtual_pipeline_model_parallel_size == 0:
+                    raise ValueError(
+                        f'num_layers: {self.num_layers} must be divisible by '
+                        f'virtual_model_parallel_size {self.virtual_pipeline_model_parallel_size}'
+                    )
+
+        if self.apply_query_key_layer_scaling:
+            self.attention_softmax_in_fp32 = True
+
+        if self.bias_activation_fusion:
+            if self.activation_func not in [F.gelu, F.silu]:
+                raise ValueError(
+                    "When bias_activation_fusion is True, activation function should be either "
+                    "gelu or swiglu"
+                )
+            if (
+                self.activation_func == F.gelu
+                and not self.gated_linear_unit
+                and not self.add_bias_linear
+            ):
+                raise ValueError(
+                    "When bias_activation_fusion is True, gated_linear_unit is False, "
+                    "and activation function is gelu, add_bias_linear must also be True."
+                )
+
+        if self.activation_func_fp8_input_store:
+            if self.activation_func != F.silu or not self.gated_linear_unit:
+                raise ValueError("Storing activation input in FP8 is supported only for SwiGLU.")
+
+        if self.apply_rope_fusion:
+            if self.rotary_interleaved:
+                raise ValueError("rotary_interleaved does not work with apply_rope_fusion.")
+
+            from megatron.core.models.common.embeddings.rope_utils import (
+                fused_apply_rotary_pos_emb,
+                fused_apply_rotary_pos_emb_thd,
+            )
+
+            if fused_apply_rotary_pos_emb is None and fused_apply_rotary_pos_emb_thd is None:
+                raise ValueError(
+                    "apply_rope_fusion is not available. Please install TE >= 1.4 or Apex."
+                )
+
+        if self.multi_latent_attention and self.rotary_interleaved:
+            raise ValueError("rotary_interleaved does not work with multi_latent_attention.")
+
+        if self.init_method is None:
+            self.init_method = init_method_normal(self.init_method_std)
+
+        if self.output_layer_init_method is None:
+            self.output_layer_init_method = scaled_init_method_normal(
+                self.init_method_std, self.num_layers
+            )
+
+        if (
+            self.moe_token_dispatcher_type == "alltoall_seq"
+            and self.tensor_model_parallel_size != self.expert_tensor_parallel_size
+        ):
+            raise ValueError(
+                "alltoall_seq dispatcher not support different TP size for MoE and Dense layer."
+            )
+
+        if self.num_moe_experts and self.fp8:
+            # TE version below 1.7.0 will raise Error when handle zeros tokens for expert
+            if not is_te_min_version("1.7.0.dev0"):
+                raise ValueError(
+                    "Only transformer-engine>=1.7.0 supports MoE FP8 training, "
+                    f"but your version is {get_te_version()}."
+                )
+
+            if self.moe_grouped_gemm and not is_te_min_version("1.11.0"):
+                raise ValueError(
+                    "Only transformer-engine>=1.11.0 supports FP8 grouped gemm, "
+                    f"but your version is {get_te_version()}."
+                )
+
+        if self.moe_router_topk_limited_devices:
+            if self.moe_router_topk_limited_devices > self.expert_model_parallel_size:
+                raise ValueError(
+                    f"moe_router_topk_limited_devices: {self.moe_router_topk_limited_devices} "
+                    f"must be smaller than expert_model_parallel_size "
+                    f"{self.expert_model_parallel_size}"
+                )
+
+        if self.flash_decode and self.fp8:
+            raise ValueError("FP8 inference is currently not support with flash decoding.")
+
+        if self.enable_cuda_graph:
+            if self.cpu_offloading:
+                raise ValueError("CUDA graphs not supported with CPU offloading.")
+            if self.recompute_granularity:
+                raise ValueError("CUDA graphs not supported with activation recomputation.")
+
+        if self.moe_token_dispatcher_type in ['allgather', 'alltoall_seq']:
+            if self.variable_seq_lengths is True:
+                raise ValueError(
+                    f"Token dispatcher type: {self.moe_token_dispatcher_type} does not support "
+                    f"variable sequence length, please use alltoall dispatcher instead."
+                )
+
+        if self.cp_comm_type is not None:
+            if isinstance(self.cp_comm_type, list):
+                assert len(self.cp_comm_type) == self.num_layers, (
+                    f"Length of cp_comm_type ({len(self.cp_comm_type)}) should equal to "
+                    f"the total number of transformer layers ({self.num_layers})!"
+                )
+            else:
+                assert isinstance(
+                    self.cp_comm_type, str
+                ), "Unsupported communication type for context parallelism!"
+
+
+@dataclass
+class MLATransformerConfig(TransformerConfig):
+    """Configuration object for megatron-core Multi-Latent Attention (MLA) transformers.
+
+    The initialization function has an argument for each parameter, including those in
+    ModelParallelConfig. Included YaRN RoPE parameters that is fused in MLA.
+    """
+
+    multi_latent_attention: bool = True
+    """Whether to use Multi-Latent Attention."""
+
+    q_lora_rank: int = 512
+    """Rank of Query tensor's low rank representation."""
+
+    kv_lora_rank: int = 512
+    """Rank of Key and Value tensors' low rank representation."""
+
+    qk_head_dim: int = 128
+    """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim"""
+
+    qk_pos_emb_head_dim: int = 64
+    """Dimension of the position embedding in the QK projection."""
+
+    v_head_dim: int = 128
+    """Dimension of the head in the V projection."""
+
+    rotary_base: float = 10000
+    """Rotary base for the rotary embeddings."""
+
+    rotary_scaling_factor: float = 40
+    """Rotary scaling factor for the rotary embeddings."""
+
+    normalization: str = "RMSNorm"
+    """Default normalization layer for MLA models is RMSNorm."""
+
+    max_position_embeddings: int = 163840
+    """Maximum position embeddings for the original model."""
+
+    beta_fast: float = 32
+    """Beta fast for YaRN RoPE."""
+
+    beta_slow: float = 1
+    """Beta slow for YaRN RoPE."""
+
+    mscale: float = 0.707
+    """Mscale for YaRN RoPE in Multi-Latent Attention."""
+
+    mscale_all_dim: float = 0.707
+    """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention."""
+
+
+import megatron.core.transformer.transformer_config
+
+megatron.core.transformer.transformer_config.TransformerConfig = TransformerConfig
+megatron.core.transformer.transformer_config.MLATransformerConfig = MLATransformerConfig
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/__init__.py b/megatron-lm-musa-patch/musa_patch/__init__.py
new file mode 100644
index 000000000..14718c0f4
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/__init__.py
@@ -0,0 +1,226 @@
+import os
+import sys
+import torch
+import torch.utils
+import torch.utils.data
+import torch_musa
+from contextlib import nullcontext
+
+def patch_before_import_megatron():
+    # Import fused_layer_norm before transformer_engine
+    from . import fused_layer_norm
+
+    from transformer_engine.pytorch.utils import get_device_compute_capability
+    def _get_device_compute_capability():
+        return (8, 0)
+    get_device_compute_capability = _get_device_compute_capability
+    from packaging.version import Version as PkgVersion
+    from transformer_engine.pytorch.attention import _flash_attn_version
+    _flash_attn_version = PkgVersion("2.5.0")
+    # Import other necessary modules to patch
+    # from . import transformer_config
+    from . import dot_product_attention
+    from . import checkpointing
+    from . import training
+    from . import linear_with_grad_accumulation_and_async_allreduce
+    from . import rotary_pos_embedding
+    from . import p2p_communication
+    from . import fused_bias_swiglu
+    if int(os.getenv("USE_MUSA_MOE", 0)):
+        from . import moe_utils
+        from . import router
+    from . import arguments
+    if int(os.getenv("USE_RECOMPUTE_VARIANCE", 0)):
+        from . import recomupte_variance
+    if int(os.getenv("USE_EPX", 0)):
+        from . import fault_tolerance_epx
+        from . import parallel_state
+    from . import optimizer
+    if int(os.getenv("ENABLE_D2H_IN_PERMUTATION", 0)):
+        from . import token_dispatcher
+
+    from . import core_pipeline_parallel_schedules
+    # Disable some unsupprted features
+    # set_jit_fusion_options
+    def set_jit_fusion_options():
+        pass
+    import megatron.training.initialize
+    megatron.training.training.set_jit_fusion_options = set_jit_fusion_options
+    megatron.training.initialize.set_jit_fusion_options = set_jit_fusion_options
+    # Disable fused_kernels
+    import megatron.legacy.fused_kernels
+    megatron.legacy.fused_kernels.load = lambda args : None
+    # Disable _compile_dependencies
+    def _compile_dependencies():
+        pass
+    megatron.training.initialize._compile_dependencies = _compile_dependencies
+
+
+def patch_after_import_torch():
+    # 1. Patch for torch.xxx
+    torch.cuda.is_available = torch.musa.is_available
+    torch.cuda.current_device = lambda : f'musa:{torch.musa.current_device()}'
+    torch.cuda.device_count = torch.musa.device_count
+    torch.cuda.set_device = torch.musa.set_device
+    torch.cuda.DoubleTensor = torch.musa.DoubleTensor
+    torch.cuda.FloatTensor = torch.musa.FloatTensor
+    torch.cuda.LongTensor = torch.musa.LongTensor
+    torch.cuda.HalfTensor = torch.musa.HalfTensor
+    torch.cuda.BFloat16Tensor = torch.musa.BFloat16Tensor
+    torch.cuda.IntTensor = torch.musa.IntTensor
+    torch.cuda.synchronize = torch.musa.synchronize
+    torch.cuda.get_rng_state = torch.musa.get_rng_state
+    torch.cuda.set_rng_state = torch.musa.set_rng_state
+    torch.cuda.random.get_rng_state = torch.musa.get_rng_state
+    torch.cuda.synchronize = torch.musa.synchronize
+    torch.cuda.empty_cache = torch.musa.empty_cache
+    torch.Tensor.cuda = torch.Tensor.musa
+    torch.cuda.manual_seed = torch.musa.manual_seed
+    torch.cuda.Event = torch.musa.Event
+    torch.cuda.Stream = torch.musa.Stream
+    torch.cuda.stream = torch.musa.stream
+    torch.cuda.get_device_properties = torch.musa.get_device_properties
+    # add torch.musa.current_devce() to activate torch.musa.default_generators
+    d = torch.musa.current_device()
+    torch.cuda.default_generators = torch.musa.default_generators
+    # torch.cuda.amp = torch.musa.amp
+    # Memory
+    torch.cuda.memory_allocated = torch.musa.memory_allocated
+    torch.cuda.max_memory_allocated = torch.musa.max_memory_allocated
+    torch.cuda.memory_reserved = torch.musa.memory_reserved
+    torch.cuda.max_memory_reserved = torch.musa.max_memory_reserved
+    torch.cuda.memory._record_memory_history = torch.musa.memory._record_memory_history
+    torch.cuda.memory._snapshot = torch.musa.memory._snapshot
+
+    # (yehua.zhang) replace lazy_call to avoid cpu memory leak,
+    # because failure of cuda init in lazy_call will cause endless operation of emplace back.
+    torch.cuda._lazy_call = torch.musa.core._lazy_init._lazy_call
+    torch.cuda._lazy_init = torch.musa.core._lazy_init._lazy_init
+
+    # 2.Patch for torch args related to cuda/musa
+    def hook_cuda_device(device):
+        if isinstance(device, str) and device.startswith("cuda"):
+            return device.replace("cuda", "musa")
+        if isinstance(device, torch.device) and device.type == "cuda":
+            return torch.device("musa", device.index)
+        return device
+
+    def maybe_hook_cuda_args(args, kwargs):
+        new_args = []
+        for arg in args:
+            new_args.append(hook_cuda_device(arg))
+        if "device" in kwargs:
+            v = kwargs["device"]
+            kwargs['device'] = hook_cuda_device(v)
+        return tuple(new_args), kwargs
+
+    # retain torch.full reference
+    original_full = torch.full
+    # redeine torch.zeros
+    def patched_full(*args, **kwargs):
+        args, kwargs = maybe_hook_cuda_args(args, kwargs)
+        result = original_full(*args, **kwargs)
+        return result
+    torch.full = patched_full
+
+    # retain torch.tensor reference
+    original_tensor = torch.tensor
+    # redefine torch.tensor
+    def patched_tensor(*args, **kwargs):
+        if 'device' in kwargs and kwargs['device'] == 'cuda':
+            kwargs['device'] = 'musa'
+        result = original_tensor(*args, **kwargs)
+        return result
+    torch.tensor = patched_tensor
+
+    # redefine torch.Tensor type
+    orig_type = torch.Tensor.type
+    def musa_type(*args, **kwargs):
+        result = orig_type(*args, **kwargs)
+        if isinstance(result, str):
+            result = result.replace("musa", "cuda")
+        return result
+    torch.Tensor.type = musa_type
+
+    # retain torch.zeros reference
+    original_zeros = torch.zeros
+    # redeine torch.zeros
+    def patched_zeros(*args, **kwargs):
+        if 'device' in kwargs and kwargs['device'] == 'cuda':
+            kwargs['device'] = 'musa'
+        result = original_zeros(*args, **kwargs)
+        return result
+    torch.zeros = patched_zeros
+
+    # retain torch.empty reference
+    original_empty = torch.empty
+    # redifine torch.empty
+    def patched_empty(*args, **kwargs):
+        if 'device' in kwargs and kwargs['device'] == 'cuda':
+            kwargs['device'] = 'musa'
+        result = original_empty(*args, **kwargs)
+        return result
+    torch.empty = patched_empty
+
+    # Original tensor class
+    original_is_cuda = torch.Tensor.is_cuda
+    def always_cuda(self):
+        return self.is_musa
+
+    # TODO : this patch may be override by transformer_engine patch
+    # we'd better unify this patch with transformer_engine patch.
+    torch.Tensor.is_cuda = property(always_cuda)
+
+    # 3. Patch for nccl/mccl
+    origin_init_process_group = torch.distributed.init_process_group
+    def patched_init_process_group(*args, **kwargs):
+        if 'backend' in kwargs and kwargs['backend'] == 'nccl':
+            kwargs['backend'] = 'mccl'
+        result = origin_init_process_group(*args, **kwargs)
+        return result
+    torch.distributed.init_process_group = patched_init_process_group
+
+    # 3. disable pin memory
+    # def pin_memory(data, device=None):
+    #     return data
+    # torch.utils.data._utils.pin_memory.pin_memory = pin_memory
+
+    # 4. disable nvtx
+    def _pass_pvtx(*args, **kwargs):
+        return
+    torch.cuda.nvtx.range_push = _pass_pvtx
+    torch.cuda.nvtx.range_pop = _pass_pvtx
+
+    # 5. disable dynamo
+    import os
+    os.environ["NVTE_TORCH_COMPILE"] = "0"
+    os.environ["TORCHDYNAMO_DISABLE"] = "1"
+
+    def noop(func):
+        return func
+    torch.compile = noop
+    torch.jit.script = noop
+
+    def get_device_capability_musa():
+        return [8, 3]
+    torch.cuda.get_device_capability = get_device_capability_musa
+
+def py_patch():
+    if sys.version_info >= (3.9, 0):
+        return
+    import math
+    def lcm(a, b):
+        return abs(a * b) // math.gcd(a, b)
+    math.lcm = lcm
+    return
+
+# Apply patch
+py_patch()
+
+patch_after_import_torch()
+
+if os.getenv("ENABLE_ZERO_BUBBLE", "0") == "1":
+    from .import zbb_light
+    zbb_light.patch_megatron()
+
+patch_before_import_megatron()
diff --git a/megatron-lm-musa-patch/musa_patch/arguments.py b/megatron-lm-musa-patch/musa_patch/arguments.py
new file mode 100644
index 000000000..e42970aba
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/arguments.py
@@ -0,0 +1,226 @@
+import dataclasses
+import torch
+import torch.nn.functional as F
+
+import megatron.training.arguments
+from megatron.training.activations import squared_relu
+from megatron.core.transformer.transformer_config import TransformerConfig, MLATransformerConfig
+moe_freq_type = megatron.training.arguments.moe_freq_type
+
+def _add_moe_args(parser):
+    group = parser.add_argument_group(title="moe")
+    # General arguments
+    group.add_argument('--expert-model-parallel-size', type=int, default=1,
+                       help='Degree of expert model parallelism.')
+    group.add_argument('--expert-tensor-parallel-size', type=int, default=None,
+                       help='Degree of expert model parallelism. Default is None, which will be set to the value of --tensor-model-paralle-size.')
+    group.add_argument('--num-experts', type=int, default=None,
+                       help='Number of Experts in MoE (None means no MoE)')
+    group.add_argument('--moe-layer-freq', type=moe_freq_type, default=1,
+                       help='Frequency between MoE layers and Dense layers. Accepts either: '
+                            '- An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers '
+                            '- A string containing a Python list expression that defines a custom pattern, e.g.: '
+                            '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] '
+                            'where 1 indicates an expert layer and 0 indicates a dense layer. '
+                            'Examples: "([0]+[1]*23)": 1 dense layer followed by 23 experts layers, '
+                            '"([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice.')
+    group.add_argument('--moe-ffn-hidden-size', type=int, default=None,
+                       help='The hidden size of each expert\'s feed-forward network (ffn). '
+                       'If not specified, defaults to the ffn_hidden_size.')
+    group.add_argument('--moe-shared-expert-intermediate-size', type=int, default=None,
+                       help='Shared expert total ffn hidden size. '
+                       'It should be equal to "num_shared_experts * ffn_size_of_each_shared_expert" if there are multiple shared experts. '
+                       'None means no shared expert.')
+    group.add_argument('--moe-shared-expert-overlap', action='store_true',
+                       help='Enable overlapping between shared expert computations and dispatcher communications. '
+                       'Without this, the shared epxerts execute after the routed experts. '
+                       'Only effective when moe-shared-expert-intermediate-size is set.')
+    group.add_argument('--moe-grouped-gemm', action='store_true',
+                       help='When there are multiple experts per rank, launch multiple local GEMM kernels in multiple streams to improve the utilization and performance with GroupedLinear in TransformerEngine.')
+    # Router arguments
+    group.add_argument('--moe-router-load-balancing-type', type=str,
+                       choices=['aux_loss', 'seq_aux_loss', 'sinkhorn', 'none'],
+                       default='aux_loss',
+                       help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".')
+    group.add_argument('--moe-router-score-function', type=str,
+                       choices=['softmax', 'sigmoid'],
+                       default='softmax',
+                       help='Score function for MoE TopK routing. Can be "softmax" or "sigmoid".')
+    group.add_argument('--moe-router-topk', type=int, default=2,
+                       help='Number of experts to route to for each token. The default is 2.')
+    group.add_argument('--moe-router-pre-softmax', action='store_true',
+                       help='Enable pre-softmax routing for MoE, which means softmax is before the top-k selection. By default, softmax is done after top-k.')
+    group.add_argument('--moe-router-num-groups', type=int, default=None,
+                       help='Number of groups to divide experts into for group-limited routing. When using group-limited routing: 1) Experts are divided into equal-sized groups, 2) For each token, a subset of groups are selected based on routing scores (sum of top-2 expert scores within each group), 3) From these selected groups, moe_router_topk experts are chosen.'
+                       'Two common use cases: 1) Device-limited routing: Set equal to expert parallel size (EP) to limit each token to experts on a subset of devices (See DeepSeek-V2: https://arxiv.org/pdf/2405.04434) 2) Node-limited routing: Set equal to number of nodes in EP group to limit each token to experts on a subset of nodes (See DeepSeek-V3: https://arxiv.org/pdf/2412.19437)')
+    group.add_argument('--moe-router-group-topk', type=int, default=None,
+                       help='Number of selected groups for group-limited routing.')
+    group.add_argument('--moe-router-topk-scaling-factor', type=float, default=None,
+                       help='Scaling factor for routing score in top-k selection, only works when --moe-router-pre-softmax enabled. Defaults to None, which means no scaling.')
+    group.add_argument('--moe-router-enable-expert-bias', action='store_true',
+                       help='TopK routing with dynamic expert bias in the aux-loss-free load balancing strategy. '
+                       'The routing decision is based on the sum of the routing scores and the expert bias. '
+                       'See https://arxiv.org/abs/2408.15664 for details.')
+    group.add_argument('--moe-router-bias-update-rate', type=float, default=1e-3,
+                       help='Expert bias update rate in the aux-loss-free load balancing strategy. '
+                       'The expert bias is updated based on the number of assigned tokens to each expert in a global batch, '
+                       'where the bias is increased for the experts with less assigned tokens and decreased for the experts with more assigned tokens. '
+                       'The default value 1e-3 is same as that used in DeepSeekV3.')
+    group.add_argument('--moe-use-legacy-grouped-gemm', action='store_true',
+                       help='Use legacy GroupedMLP rather than TEGroupedMLP. Note: The legacy one will be deprecated soon.')
+    group.add_argument('--moe-aux-loss-coeff', type=float, default=0.0,
+                       help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.')
+    group.add_argument('--moe-z-loss-coeff', type=float, default=None,
+                       help='Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended.')
+    group.add_argument('--moe-input-jitter-eps', type=float, default=None,
+                       help='Add noise to the input tensor by applying jitter with a specified epsilon value.')
+    group.add_argument('--moe-token-dispatcher-type', type=str,
+                       choices=['allgather', 'alltoall', 'flex', 'alltoall_seq'],
+                       default='allgather',
+                       help="The type of token dispatcher to use. The default is 'allgather'. Options are 'allgather', 'alltoall' and 'alltoall_seq'. We recommend using 'alltoall' when applying expert parallelism. For more information, please refer to the documentation in core/moe/README.")
+    group.add_argument('--moe-enable-deepep', action='store_true',
+                       help='[Experimental] Enable DeepSeek/DeepEP for efficient token dispatching and combine in MoE models. Only works with flex token dispatcher by setting --moe-token-dispatcher-type=flex.')
+    group.add_argument('--moe-per-layer-logging', action='store_true',
+                       help='Enable per-layer logging for MoE, currently supports auxiliary loss and z loss.')
+    # Token dropping arguments
+    group.add_argument('--moe-expert-capacity-factor', type=float, default=None,
+                       help='The capacity factor for each expert, None means no token will be dropped.')
+    group.add_argument('--moe-pad-expert-input-to-capacity', action='store_true',
+                       help='Pads the input for each expert to match the expert capacity length, effective only after the --moe-expert-capacity-factor is set.')
+    group.add_argument('--moe-token-drop-policy', type=str, default='probs', choices=['probs', 'position'],
+                       help='The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with the lowest probabilities will be dropped. If "position", tokens at the end of each batch will be dropped.')
+    group.add_argument('--moe-layer-recompute', action='store_true',
+                       help='Enable checkpointing for moe_layer, should be used when memory is not sufficient.')
+    group.add_argument('--moe-extended-tp', action='store_true',
+                       help='Deprecated. Use --expert-tensor-parallel-size instead.')
+    group.add_argument('--moe-use-upcycling', action='store_true',
+                       help='Load a checkpoint of a dense model, convert it into an MoE model, and save the converted model to the path specified by --save. '
+                       'Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.')
+    group.add_argument('--moe-permute-fusion', action='store_true',
+                       help='Fuse token rearrangement ops during token dispatching.')
+
+    # HACK(huang.huang): control dp_reduce position: tp-only-amax-red
+    group.add_argument('--tp-only-amax-red', action='store_true',
+                        help="Whether to reduce the FP8 AMAX only in the TP or TP-CP domain")
+    ## HACK(huang.huang)
+
+    # HACK(yehua.zhang): add dsv2 & dsv3 loss, q-rms-recompute
+    # dsv2
+    group.add_argument('--moe-device-level-aux-loss-coeff', type=float, default=None,
+                       help='Scaling coefficient for the device-level aux loss')
+    group.add_argument('--moe-comm-aux-loss-coeff', type=float, default=None,
+                       help='Scaling coefficient for the communication aux loss')
+    group.add_argument('--moe-device-level-capacity', action='store_true',
+                       help='Whether to consider the expert capacity of a group together')
+
+    # dsv3
+    group.add_argument('--moe-complementary-seq-aux-loss', action='store_true',
+                       help='use complementary sequence-wise aux loss in MoE, should only used with seq_aux_loss')
+    group.add_argument('--moe-router-norm-topk-prob', action='store_true',
+                       help='Enable normalization for sigmoid score in MoE, should only used with moe-router-use-sigmoid')
+
+    # q-rms-recompute
+    group = parser.add_argument_group(title="mla")
+    group.add_argument('--q-rms-recompute', action='store_true',
+                       help="use q uproj rmsnorm recompute")
+    ## HACK(yehua.zhang)
+
+    # HACK(huang.huang): add attn-recompute, recompute-variance, groupMLP_recompute
+    group.add_argument('--attn-recompute', action='store_true',
+                       help="use attn recompute")
+    group.add_argument('--mla-rms-recompute', action='store_true',
+                       help="use rms recompute before mla")
+    group.add_argument('--mlp-rms-recompute', action='store_true',
+                       help="use rms recompute before mlp")
+    group.add_argument('--recompute-variance', action='store_true',
+                       help="use recompute variance")
+    group.add_argument('--mlp-recompute', action='store_true',
+                       help="use groupMLP_recompute to recompute groupgemm and shared_exp in moelayer, mlp in dense")
+    ## HACK(huang.huang)
+    return parser
+
+
+def core_transformer_config_from_args(args, config_class=None):
+
+    # Config class.
+    config_class = config_class or TransformerConfig
+
+    if args.multi_latent_attention:
+        config_class = MLATransformerConfig
+
+    # Translate args to core transformer configuration
+    kw_args = {}
+    for f in dataclasses.fields(config_class):
+        if hasattr(args, f.name):
+            kw_args[f.name] = getattr(args, f.name)
+    kw_args['persist_layer_norm'] = not args.no_persist_layer_norm
+    kw_args['layernorm_zero_centered_gamma'] = args.apply_layernorm_1p
+    kw_args['layernorm_epsilon'] = args.norm_epsilon
+    kw_args['deallocate_pipeline_outputs'] = True
+    kw_args['pipeline_dtype'] = args.params_dtype
+    kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm
+    kw_args['num_moe_experts'] = args.num_experts
+    kw_args['rotary_interleaved'] = args.rotary_interleaved
+    kw_args['num_layers_in_first_pipeline_stage']= args.decoder_first_pipeline_num_layers
+    kw_args['num_layers_in_last_pipeline_stage']= args.decoder_last_pipeline_num_layers
+    if args.swiglu:
+        kw_args['activation_func'] = F.silu
+        kw_args['gated_linear_unit'] = True
+        kw_args['bias_activation_fusion'] = args.bias_swiglu_fusion
+    else:
+        kw_args['bias_activation_fusion'] = args.bias_gelu_fusion
+    if args.squared_relu:
+        assert not args.swiglu
+        kw_args['activation_func'] = squared_relu
+    if args.init_method_xavier_uniform:
+        kw_args['init_method'] = torch.nn.init.xavier_uniform_
+        kw_args['scaled_init_method'] = torch.nn.init.xavier_uniform_
+    if args.group_query_attention:
+        kw_args['num_query_groups'] = args.num_query_groups
+    else:
+        kw_args['num_query_groups'] = None
+    kw_args['config_logger_dir'] = args.config_logger_dir
+
+    if len(args.cp_comm_type) == 1:
+        kw_args['cp_comm_type'] = args.cp_comm_type[0]
+
+    # Return config.
+
+    # HACK(yehua.zhang): add dsv2 & dsv3 loss, mtp, q-rms-recompute from args to transformer config
+    config_instance = config_class(**kw_args)
+
+    config_instance.moe_device_level_aux_loss_coeff = args.moe_device_level_aux_loss_coeff
+    config_instance.moe_comm_aux_loss_coeff = args.moe_comm_aux_loss_coeff
+    config_instance.moe_device_level_capacity = args.moe_device_level_capacity
+
+    config_instance.moe_complementary_seq_aux_loss = args.moe_complementary_seq_aux_loss
+    config_instance.moe_router_norm_topk_prob = args.moe_router_norm_topk_prob
+    config_instance.moe_device_level_capacity = args.moe_device_level_capacity
+
+    config_instance.q_rms_recompute = args.q_rms_recompute
+    ## HACK(yehua.zhang)
+
+    # HACK(huang.huang): add attn-recompute, recompute-variance, mlp_recompute
+    config_instance.attn_recompute = args.attn_recompute
+    config_instance.mla_rms_recompute = args.mla_rms_recompute
+    config_instance.mlp_rms_recompute = args.mlp_rms_recompute
+    config_instance.recompute_variance = args.recompute_variance
+    config_instance.mlp_recompute = args.mlp_recompute
+    ## HACK(huang.huang)
+
+    # HACK(huang.huang): args check for pp=1 and first/last stage num layer=None
+    if config_instance.pipeline_model_parallel_size == 1:
+        assert config_instance.num_layers_in_first_pipeline_stage is None and config_instance.num_layers_in_last_pipeline_stage is None, \
+            f"For pipeline_model_parallel_size=1, first/last must be None, but get {config_instance.num_layers_in_first_pipeline_stage}/{config_instance.num_layers_in_last_pipeline_stage}"
+    ## HACK(huang.huang)
+
+    # HACK(huang.huang): control dp_reduce position: tp-only-amax-red
+    config_instance.tp_only_amax_red = args.tp_only_amax_red
+    ##HACK(huang.huang)
+
+    print('config_instance is ', config_instance)
+    return config_instance
+
+
+megatron.training.arguments._add_moe_args = _add_moe_args
+megatron.training.arguments.core_transformer_config_from_args = core_transformer_config_from_args
diff --git a/megatron-lm-musa-patch/musa_patch/checkpointing.py b/megatron-lm-musa-patch/musa_patch/checkpointing.py
new file mode 100644
index 000000000..45d08120f
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/checkpointing.py
@@ -0,0 +1,76 @@
+import os
+
+from megatron.training.global_vars import (
+    get_args,
+)
+
+def save_checkpoint(iteration, model, optimizer, opt_param_scheduler,
+                    num_floating_point_operations_so_far,
+                    checkpointing_context=None, pipeline_rank=None,
+                    expert_rank=None, tensor_rank=None,
+                    pipeline_parallel=None, expert_parallel=None,
+                    non_persistent_ckpt=False,
+                    train_data_iterator=None, ft_client=None,
+                    preprocess_common_state_dict_fn=None
+                    ):
+  try:
+    from dlrover.trainer.torch.flash_checkpoint.megatron_dist_ckpt \
+      import save_checkpoint as dlrover_save_checkpoint_dist
+    from dlrover.trainer.torch.flash_checkpoint.megatron \
+      import save_checkpoint as dlrover_save_checkpoint
+  except Exception as e:
+    print(f"import flash_ckpt failed {str(e)}")
+    return
+
+  args = get_args()
+  if args.use_distributed_optimizer and not args.no_save_optim:
+    dlrover_save_checkpoint_dist(iteration, model, optimizer,
+                                 opt_param_scheduler, 0,
+                                 preprocess_common_state_dict_fn)
+  else:
+    dlrover_save_checkpoint(iteration, model, optimizer,
+                            opt_param_scheduler, 0)
+
+def load_checkpoint(model, optimizer, opt_param_scheduler,
+                    load_arg='load', strict=True,
+                    ft_client=None, checkpointing_context=None,
+                    skip_load_to_model_and_opt=False):
+  try:
+    from dlrover.trainer.torch.flash_checkpoint.megatron_dist_ckpt \
+      import load_checkpoint as dlrover_load_checkpoint_dist
+    from dlrover.trainer.torch.flash_checkpoint.megatron \
+      import load_checkpoint as dlrover_load_checkpoint
+  except Exception as e:
+    print(f"import flash_ckpt failed {str(e)}")
+    return 0, 0
+
+  i = 0
+  args = get_args()
+  if args.use_distributed_optimizer and not args.no_save_optim:
+    i,  num_floating_point_operations_so_far = dlrover_load_checkpoint_dist(model,
+                                        optimizer,
+                                        opt_param_scheduler,
+                                        load_arg,
+                                        strict)
+  else:
+    i, num_floating_point_operations_so_far = dlrover_load_checkpoint(model,
+                                optimizer,
+                                opt_param_scheduler,
+                                load_arg,
+                                strict,
+                                ft_client=ft_client,
+                                checkpointing_context=checkpointing_context,
+                                skip_load_to_model_and_opt=skip_load_to_model_and_opt)
+
+  return i, num_floating_point_operations_so_far
+
+enable_async_ckpt = int(os.getenv("ENABLE_ASYNC_CKPT", 0))
+if enable_async_ckpt:
+  print("flash ckpt enabled")
+  import megatron.training.checkpointing
+
+  megatron.training.checkpointing.save_checkpoint = save_checkpoint
+  megatron.training.checkpointing.load_checkpoint = load_checkpoint
+
+  megatron.training.training.save_checkpoint = save_checkpoint
+  megatron.training.training.load_checkpoint = load_checkpoint
diff --git a/megatron-lm-musa-patch/musa_patch/core_pipeline_parallel_schedules.py b/megatron-lm-musa-patch/musa_patch/core_pipeline_parallel_schedules.py
new file mode 100644
index 000000000..a0b6e9f1a
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/core_pipeline_parallel_schedules.py
@@ -0,0 +1,61 @@
+import os
+import megatron
+import functools
+import torch
+
+
+def record_function_decorator(func):
+    @functools.wraps(func)
+    def new_func(*args, **kwargs):
+        with torch.profiler.record_function(func.__name__):
+            return func(*args, **kwargs)
+
+    return new_func
+
+
+original_forward_step = megatron.core.pipeline_parallel.schedules.forward_step
+original_backward_step = megatron.core.pipeline_parallel.schedules.backward_step
+
+
+@record_function_decorator
+def forward_step(
+    forward_step_func,
+    data_iterator,
+    model,
+    num_microbatches,
+    input_tensor,
+    forward_data_store,
+    config,
+    collect_non_loss_data=False,
+    checkpoint_activations_microbatch=None,
+    is_first_microbatch=False,
+    current_microbatch=None,
+    encoder_decoder_xattn=False,
+):
+    return original_forward_step(
+        forward_step_func,
+        data_iterator,
+        model,
+        num_microbatches,
+        input_tensor,
+        forward_data_store,
+        config,
+        collect_non_loss_data,
+        checkpoint_activations_microbatch,
+        is_first_microbatch,
+        current_microbatch,
+        encoder_decoder_xattn,
+    )
+
+
+@record_function_decorator
+def backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config):
+    return original_backward_step(
+        input_tensor, output_tensor, output_tensor_grad, model_type, config
+    )
+
+
+enable_profiler = int(os.getenv("ENABLE_PROFILER", 0))
+if enable_profiler:
+    megatron.core.pipeline_parallel.schedules.forward_step = forward_step
+    megatron.core.pipeline_parallel.schedules.backward_step = backward_step
diff --git a/megatron-lm-musa-patch/musa_patch/dot_product_attention.py b/megatron-lm-musa-patch/musa_patch/dot_product_attention.py
new file mode 100644
index 000000000..6cd629b05
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/dot_product_attention.py
@@ -0,0 +1,262 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+
+import math
+
+import torch
+from torch import Tensor
+
+from megatron.core import parallel_state, tensor_parallel
+from megatron.core.fusions.fused_softmax import FusedScaleMaskSoftmax
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import attention_mask_func
+from megatron.core.utils import divide
+
+
+
+class DotProductAttention(MegatronModule):
+    """
+    Region where selective activation recomputation is applied.
+    This region is memory intensive but less compute intensive which
+    makes activation checkpointing more efficient for LLMs (20B+).
+    See Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198 for more details.
+
+    We use the following notation:
+     h: hidden size
+     n: number of attention heads
+     p: number of tensor model parallel partitions
+     b: batch size
+     s: sequence length
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        layer_number: int,
+        attn_mask_type: AttnMaskType,
+        attention_type: str,
+        attention_dropout: float = None,
+        softmax_scale: float = None,
+        k_channels: int = None,
+        v_channels: int = None,
+        cp_comm_type: str = None,
+    ):
+        super().__init__(config=config)
+
+        self.config: TransformerConfig = config
+
+        assert (
+            self.config.context_parallel_size == 1
+        ), "Context parallelism is only supported by TEDotProductAttention!"
+
+        assert (
+            self.config.window_size is None
+        ), "Sliding Window Attention is only supported by TEDotProductAttention!"
+        print("use DotProductAttention")
+        self.layer_number = max(1, layer_number)
+        self.attn_mask_type = attn_mask_type
+        self.attention_type = attention_type  # unused for now
+
+        projection_size = self.config.kv_channels * self.config.num_attention_heads
+        from megatron.training import get_args
+        args = get_args()
+        self.use_flash_attn = args.use_flash_attn and not args.multi_latent_attention
+
+        # Per attention head and per partition values.
+        world_size = parallel_state.get_tensor_model_parallel_world_size()
+        self.hidden_size_per_partition = divide(projection_size, world_size)
+        self.hidden_size_per_attention_head = divide(projection_size, config.num_attention_heads)
+        self.num_attention_heads_per_partition = divide(self.config.num_attention_heads, world_size)
+        self.num_query_groups_per_partition = divide(self.config.num_query_groups, world_size)
+
+        coeff = None
+        self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
+        if self.config.apply_query_key_layer_scaling:
+            coeff = self.layer_number
+            self.norm_factor *= coeff
+
+        self.scale_mask_softmax = FusedScaleMaskSoftmax(
+            input_in_fp16=self.config.fp16,
+            input_in_bf16=self.config.bf16,
+            attn_mask_type=self.attn_mask_type,
+            scaled_masked_softmax_fusion=self.config.masked_softmax_fusion,
+            mask_func=attention_mask_func,
+            softmax_in_fp32=self.config.attention_softmax_in_fp32,
+            scale=coeff,
+        )
+
+        # Dropout. Note that for a single iteration, this layer will generate
+        # different outputs on different number of parallel partitions but
+        # on average it should not be partition dependent.
+        self.attention_dropout = torch.nn.Dropout(
+            self.config.attention_dropout if attention_dropout is None else attention_dropout
+        )
+
+    def _flash_attn_impl(self, query_states, key_states, value_states, attention_mask=None):
+        # attention_mask useless for now
+        output_size = (
+            query_states.size(0),
+            query_states.size(1),
+            query_states.size(2) * query_states.size(3)
+        ) #seq_len, batch_size, head_num * head_dim
+
+        query_states = query_states.permute(1, 2, 0, 3)
+        key_states = key_states.permute(1, 2, 0, 3)
+        value_states = value_states.permute(1, 2, 0, 3)
+
+        bsz, num_heads, q_len, head_dim = query_states.size()
+        # kv_seq_len = key_states.size(2)
+        with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False):
+            attn_output = torch.nn.functional.scaled_dot_product_attention(
+                query_states, # batch_size, head_num , seq_len,  head_size
+                key_states, # batch_size, head_num, seq_len, head_size
+                value_states, #batch_size, head_num, seq_len, head_size
+                attn_mask=None,# bsz * num_heads, q_len, kv_seq_len
+                dropout_p=0.0,
+                is_causal=True,
+            )
+        if attn_output.size() != (bsz, num_heads, q_len, head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, num_heads, q_len, head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+        attn_output = attn_output.permute(2, 0, 1, 3).reshape(*output_size).contiguous()
+        return attn_output
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        attention_mask: Tensor,
+        attn_mask_type: AttnMaskType = None,
+        attention_bias: Tensor = None,
+        packed_seq_params: PackedSeqParams = None,
+    ):
+        assert packed_seq_params is None, (
+            "Packed sequence is not supported by DotProductAttention."
+            "Please use TEDotProductAttention instead."
+        )
+        # ===================================
+        # Raw attention scores. [b, n/p, s, s]
+        # ===================================
+
+        # expand the key and value [sk, b, ng, hn] -> [sk, b, np, hn]
+        # This is a noop for normal attention where ng == np. When using group query attention this
+        # creates a view that has the keys and values virtually repeated along their dimension to
+        # match the number of queries.
+
+        # attn_mask_type is not used.
+        if self.num_attention_heads_per_partition // self.num_query_groups_per_partition > 1:
+            key = key.repeat_interleave(
+                self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
+            )
+            value = value.repeat_interleave(
+                self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
+            )
+
+
+        if self.use_flash_attn:
+            # TODO(optimize mask method)
+            # if attention_mask.dtype == torch.bool:
+            #     attention_mask.mas
+            assert attn_mask_type == AttnMaskType.causal
+
+            context = self._flash_attn_impl(query, key, value, attention_mask)
+            # new_context_shape = context.size()[:-2] + (self.hidden_size_per_partition,)
+            # context = context.view(*new_context_shape)
+            return context
+        # [b, np, sq, sk]
+        output_size = (
+            query.size(1),
+            query.size(2),
+            query.size(0),
+            key.size(0),
+        )
+
+        # [sq, b, np, hn] -> [sq, b * np, hn]
+        # This will be a simple view when doing normal attention, but in group query attention
+        # the key and value tensors are repeated to match the queries so you can't use simple strides
+        # to extract the queries.
+        query = query.reshape(output_size[2], output_size[0] * output_size[1], -1)
+        # [sk, b, np, hn] -> [sk, b * np, hn]
+        key = key.view(output_size[3], output_size[0] * output_size[1], -1)
+
+
+        # preallocting input tensor: [b * np, sq, sk]
+        matmul_input_buffer = parallel_state.get_global_memory_buffer().get_tensor(
+            (output_size[0] * output_size[1], output_size[2], output_size[3]), query.dtype, "mpu",
+        )
+
+        # Raw attention scores. [b * np, sq, sk]
+        matmul_result = torch.bmm(query.transpose(0, 1), key.transpose(0, 1).transpose(1, 2))
+        matmul_result *= (1.0 / self.norm_factor)
+
+        # torch.baddbmm(
+        #     matmul_input_buffer,
+        #     query.transpose(0, 1),  # [b * np, sq, hn]
+        #     key.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]
+        #     beta=0.0,
+        #     alpha=(1.0 / self.norm_factor),
+        # )
+
+        # change view to [b, np, sq, sk]
+        attention_scores = matmul_result.view(*output_size)
+
+        # ===========================
+        # Attention probs and dropout
+        # ===========================
+
+        # attention scores and attention mask [b, np, sq, sk]
+        attention_probs: Tensor = self.scale_mask_softmax(attention_scores, attention_mask)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+
+        if not self.config.sequence_parallel:
+            with tensor_parallel.get_cuda_rng_tracker().fork():
+                attention_probs = self.attention_dropout(attention_probs)
+        else:
+            attention_probs = self.attention_dropout(attention_probs)
+
+        # =========================
+        # Context layer. [sq, b, hp]
+        # =========================
+
+        # value -> context layer.
+        # [sk, b, np, hn] --> [b, np, sq, hn]
+
+        # context layer shape: [b, np, sq, hn]
+        output_size = (
+            value.size(1),
+            value.size(2),
+            query.size(0),
+            value.size(3),
+        )
+
+        # change view [sk, b * np, hn]
+        value = value.view(value.size(0), output_size[0] * output_size[1], -1)
+
+        # change view [b * np, sq, sk]
+        attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
+
+        # matmul: [b * np, sq, hn]
+        context = torch.bmm(attention_probs, value.transpose(0, 1))
+
+        # change view [b, np, sq, hn]
+        context = context.view(*output_size)
+
+        # [b, np, sq, hn] --> [sq, b, np, hn]
+        context = context.permute(2, 0, 1, 3).contiguous()
+
+        # [sq, b, np, hn] --> [sq, b, hp]
+        new_context_shape = context.size()[:-2] + (-1,)
+        context = context.view(*new_context_shape)
+
+        return context
+
+import megatron.core.transformer.dot_product_attention
+megatron.core.transformer.dot_product_attention.DotProductAttention = DotProductAttention
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/__init__.py b/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/__init__.py
new file mode 100644
index 000000000..3d1537ff0
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/__init__.py
@@ -0,0 +1,3 @@
+from . import param_and_grad_buffer
+from . import data_samplers
+from . import epx_model_and_optimizer
diff --git a/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/data_samplers.py b/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/data_samplers.py
new file mode 100644
index 000000000..fb6bebfb2
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/data_samplers.py
@@ -0,0 +1,76 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+"""Dataloaders."""
+
+import os
+import logging
+import random
+import torch
+import numpy as np
+from torch.utils.data import Dataset
+from megatron.core import mpu
+from megatron.training import get_args
+from megatron.core.utils import log_single_rank
+from megatron.legacy.data.data_samplers import MegatronPretrainingSampler
+from megatron.legacy.data.data_samplers import MegatronPretrainingRandomSampler
+
+logger = logging.getLogger(__name__)
+
+def build_pretraining_data_loader(dataset, consumed_samples):
+    """Build dataloader given an input dataset."""
+
+    if dataset is None:
+        return None
+    args = get_args()
+
+    if int(os.getenv("USE_EPX", 0)): # Fault tolerance sampler
+        from epx import EpxSampler
+        log_single_rank(logger, logging.INFO, f"Use EpxSampler")
+
+        batch_sampler = EpxSampler(
+            dataset=None,
+            total_samples=len(dataset),
+            consumed_samples=consumed_samples,
+            micro_batch_size=args.micro_batch_size,
+            data_parallel_rank=mpu.get_data_parallel_rank(),
+            data_parallel_size=mpu.get_data_parallel_world_size())
+    else: # Megatron sampler
+        if args.dataloader_type == 'single':
+            batch_sampler = MegatronPretrainingSampler(
+                total_samples=len(dataset),
+                consumed_samples=consumed_samples,
+                micro_batch_size=args.micro_batch_size,
+                data_parallel_rank=mpu.get_data_parallel_rank(),
+                data_parallel_size=mpu.get_data_parallel_world_size())
+        elif args.dataloader_type == 'cyclic':
+            batch_sampler = MegatronPretrainingRandomSampler(
+                dataset,
+                total_samples=len(dataset),
+                consumed_samples=consumed_samples,
+                micro_batch_size=args.micro_batch_size,
+                data_parallel_rank=mpu.get_data_parallel_rank(),
+                data_parallel_size=mpu.get_data_parallel_world_size(),
+                data_sharding=args.data_sharding)
+        elif args.dataloader_type == "external":
+            # External dataloaders are passed through. User is expected to provide a
+            # torch-compatible dataloader and define samplers, if needed.
+            return dataset
+        else:
+            raise Exception('{} dataloader type is not supported.'.format(
+                    args.dataloader_type))
+
+    # Torch dataloader.
+    return torch.utils.data.DataLoader(dataset,
+                                       batch_sampler=batch_sampler,
+                                       num_workers=args.num_workers,
+                                       pin_memory=True,
+                                       persistent_workers=True if args.num_workers > 0 else False,
+                                       )
+
+
+import sys
+for k in sys.modules:
+    if k.startswith('megatron'):
+        for target in ['build_pretraining_data_loader']:
+            if getattr(sys.modules[k], target, None):
+                setattr(sys.modules[k], target, build_pretraining_data_loader)
diff --git a/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/epx_model_and_optimizer.py b/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/epx_model_and_optimizer.py
new file mode 100644
index 000000000..83bbd8283
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/epx_model_and_optimizer.py
@@ -0,0 +1,112 @@
+import logging
+import wrapt
+import megatron
+import megatron.core.parallel_state as parallel_state
+from megatron.training.training import setup_model_and_optimizer
+from epx.optim import epx_wrap_optimizer_instance
+
+from megatron.core.optimizer import ChainedOptimizer
+from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer
+
+logger = logging.getLogger(__name__)
+
+# Do not use core.transformer.module.Float16Module now
+# from megatron.core.transformer.module import Float16Module
+# setattr(sys.modules["megatron.training.training"], "Float16Module", Float16Module)
+
+@wrapt.decorator
+def setup_model_and_optimizer_wrapper(wrapped, _, args, kwargs):
+    def _dump_state_dict():
+        nonlocal optimizer
+
+        state_dict = {}
+
+        if len(model) == 1:
+            state_dict['model'] =  model[0].state_dict_for_save_checkpoint()
+        else:
+            for i in range(len(model)):
+                mpu.set_virtual_pipeline_model_parallel_rank(i)
+                state_dict['model%d' % i] =  model[i].state_dict_for_save_checkpoint()
+
+        if opt_param_scheduler is not None:
+            state_dict['opt_param_scheduler'] = opt_param_scheduler.state_dict()
+
+        # Optimizer stuff.
+        if optimizer is not None and not optimizer.is_stub_optimizer:
+            state_dict["optimizer"] = []
+            if isinstance(optimizer, ChainedOptimizer):
+                for optim in optimizer.chained_optimizers:
+                    inner_state_dict = optim.optimizer.state_dict()
+                    shard_fp32_from_float16_groups = optim.shard_fp32_from_float16_groups
+                    state_dict["optimizer"].append({ "inner_state_dict" : inner_state_dict,
+                                                        "shard_fp32_from_float16_groups": shard_fp32_from_float16_groups})
+            elif isinstance(optimizer, DistributedOptimizer):
+                inner_state_dict = optimizer.optimizer.state_dict()
+                shard_fp32_from_float16_groups = optimizer.shard_fp32_from_float16_groups
+                state_dict["optimizer"] = { "inner_state_dict" : inner_state_dict,
+                                            "shard_fp32_from_float16_groups": shard_fp32_from_float16_groups}
+            else:
+                assert False, f"epx _dump_state_dict not support {optimizer} now."
+
+        return state_dict
+
+
+    def _load_state_dict(state_dict):
+        nonlocal optimizer
+        opt_param_scheduler.load_state_dict(state_dict["opt_param_scheduler"])
+
+        if len(model) == 1:
+            model[0].load_state_dict(state_dict["model"])
+        else:
+            for i in range(len(model)):
+                mpu.set_virtual_pipeline_model_parallel_rank(i)
+                model[i] = state_dict['model%d' % i]
+
+        if optimizer is not None and not optimizer.is_stub_optimizer:
+            if isinstance(optimizer, ChainedOptimizer):
+                optimizer_states = state_dict["optimizer"]
+                assert len(optimizer_states) == len(optimizer.chained_optimizers), "optimizer state size mismatch"
+                for optim, state in zip(optimizer.chained_optimizers, optimizer_states):
+                    optim.optimizer.load_state_dict(state["inner_state_dict"])
+                    _copy_shard_params(state["shard_fp32_from_float16_groups"], optim.shard_fp32_from_float16_groups)
+            elif isinstance(optimizer, DistributedOptimizer):
+                optimizer.optimizer.load_state_dict(state_dict["optimizer"]["inner_state_dict"])
+                _copy_shard_params(state_dict["optimizer"]["shard_fp32_from_float16_groups"], optimizer.shard_fp32_from_float16_groups)
+            else:
+                assert False, f"epx _load_state_dict not support {optimizer} now."
+
+
+    def _copy_shard_params(src_params, dst_params):
+        assert len(src_params) == len(dst_params), "param size mismatch"
+        for src, dst in zip(src_params, dst_params):
+            if src is None or dst is None:
+                continue
+
+            if isinstance(src, list) and isinstance(dst, list):
+                _copy_shard_params(src, dst)
+                continue
+
+            assert isinstance(src, torch.Tensor) and isinstance(dst, torch.Tensor), "param type mismatch"
+            assert src.shape == dst.shape, "param shape mismatch"
+            assert src.dtype == dst.dtype, "param dtype mismatch"
+            dst.data.copy_(src.data)
+
+    logger.info("epx wrapped setup_model_and_optimizer")
+
+    model, optimizer, opt_param_scheduler = wrapped(*args, **kwargs)
+    lcp = parallel_state.get_epx_data_parallel_lcp()
+
+    logger.info(f"epx register replica_state")
+    lcp.register_module("replica_state", _dump_state_dict, _load_state_dict)
+
+    logger.info(f"Start wrap optimizer by epx")
+
+    optimizer = epx_wrap_optimizer_instance(optimizer, lcp)
+
+    logger.info(f"Finished wrap optimizer by epx")
+
+    return model, optimizer, opt_param_scheduler
+
+wraped_setup_model_and_optimizer = setup_model_and_optimizer_wrapper(setup_model_and_optimizer)
+
+megatron.training.training.setup_model_and_optimizer = wraped_setup_model_and_optimizer
diff --git a/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/param_and_grad_buffer.py b/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/param_and_grad_buffer.py
new file mode 100644
index 000000000..5de3b90f3
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/fault_tolerance_epx/param_and_grad_buffer.py
@@ -0,0 +1,210 @@
+import logging
+import os
+
+import torch
+from contextlib import nullcontext
+from torch.distributed import _coalescing_manager
+
+from megatron.core.distributed.param_and_grad_buffer import _ParamAndGradBucketGroup
+import megatron.core.parallel_state as parallel_state
+from megatron.core.utils import is_torch_min_version
+
+logger = logging.getLogger(__name__)
+
+if is_torch_min_version("1.13.0"):
+    dist_all_gather_func = torch.distributed.all_gather_into_tensor
+    dist_reduce_scatter_func = torch.distributed.reduce_scatter_tensor
+else:
+    dist_all_gather_func = torch.distributed._all_gather_base
+    dist_reduce_scatter_func = torch.distributed._reduce_scatter_base
+
+def shard_buffer(buffer: torch.Tensor, data_parallel_world_size: int):
+    """
+    Shard buffer into data_parallel_world_size chunks of equal size.
+    """
+    assert buffer.numel() % data_parallel_world_size == 0
+    shard_size = buffer.numel() // data_parallel_world_size
+    sharded_buffer = [
+        buffer[(r * shard_size) : ((r + 1) * shard_size)] for r in range(data_parallel_world_size)
+    ]
+    return sharded_buffer
+
+def start_grad_sync(self):
+    """
+    Initiates grad sync (all-reduce or reduce-scatter) communication operations
+    for all buckets in the bucket group.
+
+    When ddp_config.overlap_grad_reduce is set to True, dispatches an asynchronous
+    communication call. When ddp_config.overlap_grad_reduce is set to False, makes
+    synchronous call.
+    """
+    assert (
+        self.grad_reduce_handle is None
+    ), 'Should not have multiple communication calls outstanding at once'
+    #print('before')
+
+    if self.ddp_config.check_for_nan_in_grad or self.ddp_config.check_for_large_grads:
+        self.check_grads(
+            check_for_nan_or_inf=self.ddp_config.check_for_nan_in_grad,
+            check_for_large=self.ddp_config.check_for_large_grads,
+        )
+
+    # gradient_scaling_factor already takes into account whether we are computing
+    # an average or sum in the data-parallel collective.
+    for bucket in self.buckets:
+        if bucket.gradient_scaling_factor != 1.0:
+            bucket.grad_data *= bucket.gradient_scaling_factor
+
+    # Decide reduce_op.
+    reduce_op = torch.distributed.ReduceOp.SUM
+    if self.ddp_config.average_in_collective:
+        reduce_op = torch.distributed.ReduceOp.AVG
+
+    # We use the following stream synchronization for the gradient reduction
+    # within and across DistOpt instances.
+
+    # Compute Stream: -------------Gradient compute-------------------
+    # Comm. Stream:   ------(wait for NCCL)-----(wait for NCCL)-------
+    # NCCL Stream:          -------RS------     -------AR------
+
+    # Use async communications only when overlap_grad_reduce is True.
+    async_op = (
+        self.ddp_config.overlap_grad_reduce
+        and self.ddp_config.num_distributed_optimizer_instances == 1
+    )
+    if (
+        self.ddp_config.num_distributed_optimizer_instances > 1
+        and self.ddp_config.overlap_grad_reduce
+    ):
+        # Assign a communication stream if we have multiple DistOpt instances and we
+        # need to overlap communication.
+        stream_context = torch.cuda.stream(self.communication_stream)
+
+        # The RS/AR communication stream needs to wait for the default stream
+        # to complete its gradient computation before launching the next
+        # gradient reduction collective.
+        self.communication_stream.wait_stream(torch.cuda.default_stream())
+    else:
+        stream_context = nullcontext()
+
+    if self.ddp_config.use_distributed_optimizer:
+        communication_group = self.intra_distributed_optimizer_instance_group
+    else:
+        communication_group = self.data_parallel_group
+
+    # Coalesce communication kernels across buckets in the bucket group.
+    with stream_context, _coalescing_manager(communication_group, async_ops=async_op) as cm:
+        for bucket in self.buckets:
+            if self.ddp_config.use_distributed_optimizer:
+                local_data_view = shard_buffer(
+                    bucket.grad_data, self.intra_distributed_optimizer_instance_size
+                )[self.intra_distributed_optimizer_instance_rank]
+
+                dist_reduce_scatter_func(
+                    local_data_view,
+                    bucket.grad_data,
+                    op=reduce_op,
+                    group=communication_group,
+                    async_op=async_op,
+                )
+
+                if int(os.getenv("USE_EPX", 0)) and not async_op:
+                    epx_sync_grad_across_instances(local_data_view)
+            else:
+                torch.distributed.all_reduce(
+                    bucket.grad_data, op=reduce_op, group=communication_group, async_op=async_op
+                )
+                if int(os.getenv("USE_EPX", 0)) and not async_op:
+                    epx_sync_grad_across_instances(bucket.grad_data)
+
+    # print('before before allreduce')
+    # With multiple DistOpt instances, we need to all-reduce across instances.
+    if (
+        self.ddp_config.use_distributed_optimizer
+        and self.ddp_config.num_distributed_optimizer_instances > 1
+    ):
+        # Create a new coalescing manager for the inter-instance all-reduce.
+        with stream_context, _coalescing_manager(
+            self.inter_distributed_optimizer_instance_group, async_ops=async_op
+        ) as cm:
+            for bucket in self.buckets:
+                local_data_view = shard_buffer(
+                    bucket.grad_data, self.intra_distributed_optimizer_instance_size
+                )[self.intra_distributed_optimizer_instance_rank]
+                # print('before all reduce')
+                torch.distributed.all_reduce(
+                    local_data_view,
+                    op=reduce_op,
+                    group=self.inter_distributed_optimizer_instance_group,
+                    async_op=async_op,
+                )
+                # print('after all reduce')
+    # print('after after allreduce')
+
+    if async_op:
+        self.grad_reduce_handle = cm
+    else:
+        # When using `_coalescing_manager`, even if a synchronous op (async_op=False) is used,
+        # `cm` is not None, which is different from when `_coalescing_manager` is not used in
+        # which case the torch.distributed._reduce_scatter_base() will return None. In order to
+        # maintain consistency with prior code, we need to manually set communication handle to
+        # None.
+        self.grad_reduce_handle = None
+
+def finish_grad_sync(self):
+    """
+    Finishes grad sync (all-reduce or reduce-scatter) communication operations
+    for all buckets in the bucket group.
+
+    When ddp_config.overlap_grad_reduce is set to True, waits for asynchronous
+    communication call to complete. When ddp_config.overlap_grad_reduce is set to False,
+    makes synchronous call.
+    """
+    self.param_gather_dispatched = False
+
+    # If overlap_grad_reduce is False, start (and finish) synchronous communication call here.
+    # print(f'before self.ddp_config.overlap_grad_reduce is {self.ddp_config.overlap_grad_reduce}')
+    if not self.ddp_config.overlap_grad_reduce:
+        self.start_grad_sync()
+        return
+    # print(f'after self.ddp_config.overlap_grad_reduce is {self.ddp_config.overlap_grad_reduce}')
+    # When using multiple DistOpt instances, we don't need to sync here as we launch
+    # communications on a separate communication stream.
+    if self.ddp_config.num_distributed_optimizer_instances > 1:
+        torch.cuda.default_stream().wait_stream(self.communication_stream)
+        return
+    assert self.grad_reduce_handle is not None, (
+        f'Communication call has not been issued for this bucket '
+        f'({len(self.params_with_grad)}/{len(self.params)} params have grad available)'
+    )
+    self.grad_reduce_handle.wait()
+    self.grad_reduce_handle = None
+
+    # TODO: Using `_coalescing_manager` to optimize code structure.
+    if int(os.getenv("USE_EPX", 0)):
+        for bucket in self.buckets:
+            if self.ddp_config.use_distributed_optimizer:
+                local_data_view = shard_buffer(
+                    bucket.grad_data, self.intra_distributed_optimizer_instance_size
+                )[self.intra_distributed_optimizer_instance_rank]
+                if int(os.getenv("USE_EPX", 0)):
+                    epx_sync_grad_across_instances(local_data_view)
+            else:
+                if int(os.getenv("USE_EPX", 0)):
+                    epx_sync_grad_across_instances(bucket.grad_data)
+
+def epx_sync_grad_across_instances(tensor):
+    """
+    Sync grad across instances.
+    """
+    lcp = parallel_state.get_epx_data_parallel_lcp()
+    # TODO: avoid assemble before each allreduce
+    lcp.assemble()
+    logger.info("start epx allreduce")
+    logger.debug(f"grad before epx allreduce : {tensor[:10]}")
+    lcp.allreduce([tensor]).wait()
+    logger.info("finished epx allreduce")
+    logger.debug(f"grad after epx allreduce : {tensor[:10]}")
+
+_ParamAndGradBucketGroup.start_grad_sync = start_grad_sync
+_ParamAndGradBucketGroup.finish_grad_sync = finish_grad_sync
diff --git a/megatron-lm-musa-patch/musa_patch/fused_bias_swiglu.py b/megatron-lm-musa-patch/musa_patch/fused_bias_swiglu.py
new file mode 100644
index 000000000..5c2c2e58f
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/fused_bias_swiglu.py
@@ -0,0 +1,27 @@
+import torch
+import torch.nn.functional as F
+
+
+class MusaSwiGLUFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, fp8_input_store):
+        ctx.save_for_backward(input)
+        ctx.fp8_input_store = fp8_input_store
+        return torch.ops.aten._fused_swiglu_forward(input)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (input, ) = ctx.saved_tensors
+        return torch.ops.aten._fused_swiglu_backward(grad_output, input), None
+
+
+import megatron.core.fusions.fused_bias_swiglu
+megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction = MusaSwiGLUFunction
+
+# import sys
+# for k in sys.modules:
+#     if k.startswith('megatron.core.fusions.fused_bias_swiglu'):
+#         for target in ['bias_swiglu_impl']:
+#             if getattr(sys.modules[k], target, None):
+#                 print(f'target is {target}')
+#                 setattr(sys.modules[k], target, bias_swiglu_impl)
diff --git a/megatron-lm-musa-patch/musa_patch/fused_layer_norm.py b/megatron-lm-musa-patch/musa_patch/fused_layer_norm.py
new file mode 100644
index 000000000..6faadd898
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/fused_layer_norm.py
@@ -0,0 +1,99 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+import numbers
+
+import torch
+import torch_musa
+from torch import Tensor
+from torch.nn import init
+from torch.nn.parameter import Parameter
+
+from megatron.core.transformer import TransformerConfig
+
+
+
+class FusedLayerNorm(torch.nn.Module):
+
+    """Layer Norm, fused into a single CUDA kernel.
+
+    Args:
+      hidden_size (int): Transformer hidden dimension.
+
+      eps (float): Epsilon added to denominator, for numerical stability.
+
+      persist_layer_norm (bool): Use persistent fused layer norm kernel.
+      This kernel supports only a set of hidden sizes. Please
+      check persist_ln_hidden_sizes if your hidden size is supported.
+
+      zero_centered_gamma (bool): Adjust LayerNorm weights such that they are
+      centered around zero. This improves numerical stability.
+
+      config (TransformerConfig): Transformer config. Include to match custom
+      layer norm interfaces.
+
+      normalization (str): Normalization type, used for Transformer Engine.
+      Must equal 'LayerNorm' here.
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        hidden_size: int,
+        eps: float = 1e-5,
+        persist_layer_norm: bool = True,
+        zero_centered_gamma: bool = False,
+        normalization: str = "LayerNorm",  # included to match TE interface
+    ):
+        super().__init__()
+        print("use FusedLayerNorm")
+
+        self.config = config
+
+        self.zero_centered_gamma = self.config.layernorm_zero_centered_gamma
+
+        if self.config.normalization == "LayerNorm":
+            self.norm_impl = torch.layer_norm
+        elif self.config.normalization == "RMSNorm":
+            self.norm_impl = torch.rms_norm
+        else:
+            raise ValueError(f'({self.config.normalization}) is not supported in FusedLayerNorm')
+
+        if isinstance(hidden_size, numbers.Integral):
+            hidden_size = (hidden_size,)
+        # self.hidden_size = torch.Size(hidden_size)
+        self.hidden_size = hidden_size
+        self.eps = eps
+        self.weight = Parameter(torch.Tensor(*hidden_size))
+        self.bias = Parameter(torch.Tensor(*hidden_size)) if self.config.normalization == "LayerNorm" else None
+        self.reset_parameters()
+        self.sequence_parallel = self.config.sequence_parallel
+
+
+        # set sequence parallelism flag on weight and bias parameters
+        setattr(self.weight, 'sequence_parallel', self.sequence_parallel)
+        if self.config.normalization == "LayerNorm":
+            setattr(self.bias, 'sequence_parallel', self.sequence_parallel)
+
+    def reset_parameters(self):
+
+        if self.zero_centered_gamma:
+            init.zeros_(self.weight)
+            if self.config.normalization == "LayerNorm":
+                init.zeros_(self.bias)
+        else:
+            init.ones_(self.weight)
+            if self.config.normalization == "LayerNorm":
+                init.zeros_(self.bias)
+
+    def forward(self, input: Tensor) -> Tensor:
+
+        weight = self.weight + 1 if self.zero_centered_gamma else self.weight
+        if self.config.normalization == "LayerNorm":
+            output = self.norm_impl(input, self.hidden_size, weight, self.bias, self.eps)
+        else:
+            output = self.norm_impl(input, self.hidden_size, weight, self.eps)
+
+        return output
+
+import megatron.core.fusions.fused_layer_norm
+megatron.core.fusions.fused_layer_norm.FusedLayerNorm = FusedLayerNorm
diff --git a/megatron-lm-musa-patch/musa_patch/linear_with_grad_accumulation_and_async_allreduce.py b/megatron-lm-musa-patch/musa_patch/linear_with_grad_accumulation_and_async_allreduce.py
new file mode 100644
index 000000000..afeac8f63
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/linear_with_grad_accumulation_and_async_allreduce.py
@@ -0,0 +1,188 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+# Parts of the code here are adapted from PyTorch
+# repo: https://github.com/pytorch/pytorch
+
+import io
+import math
+import os
+import warnings
+from typing import Any, Callable, List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+import torch.nn.init as init
+from torch.cuda.amp import custom_bwd, custom_fwd
+from torch.nn.parameter import Parameter
+
+from megatron.core.parallel_state import (
+    get_global_memory_buffer,
+    get_tensor_model_parallel_group,
+    get_tensor_model_parallel_world_size,
+)
+
+from megatron.core.utils import prepare_input_tensors_for_wgrad_compute
+import fused_weight_gradient_mlp_cuda
+
+class LinearWithGradAccumulationAndAsyncCommunication(torch.autograd.Function):
+    """See linear_with_grad_accumulation_and_async_allreduce"""
+
+    @staticmethod
+    @custom_fwd
+    def forward(
+        ctx,
+        input,
+        weight,
+        bias,
+        gradient_accumulation_fusion,
+        allreduce_dgrad,
+        sequence_parallel,
+        grad_output_buffer,
+        wgrad_deferral_limit,
+    ):
+        ctx.save_for_backward(input, weight)
+        ctx.use_bias = bias is not None
+        ctx.gradient_accumulation_fusion = gradient_accumulation_fusion
+        ctx.allreduce_dgrad = allreduce_dgrad
+        ctx.sequence_parallel = sequence_parallel
+        ctx.wgrad_deferral_limit = wgrad_deferral_limit
+        ctx.grad_output_buffer = grad_output_buffer
+
+        if sequence_parallel:
+            world_size = get_tensor_model_parallel_world_size()
+            dim_size = list(input.size())
+            dim_size[0] = dim_size[0] * world_size
+
+            all_gather_buffer = get_global_memory_buffer().get_tensor(dim_size, input.dtype, "mpu")
+            torch.distributed._all_gather_base(
+                all_gather_buffer, input, group=get_tensor_model_parallel_group()
+            )
+            total_input = all_gather_buffer
+        else:
+            total_input = input
+
+        output = torch.matmul(total_input, weight.t())
+        if bias is not None:
+            output = output + bias
+        return output
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad_output):
+        input, weight = ctx.saved_tensors
+        use_bias = ctx.use_bias
+        grad_output_buffer = ctx.grad_output_buffer
+        wgrad_deferral_limit = ctx.wgrad_deferral_limit
+
+        wgrad_compute = True
+        # if grad_output_buffer is not None:
+        #     grad_output_buffer.append(grad_output)
+        #     wgrad_compute = False
+
+        if grad_output_buffer is not None:
+            if wgrad_deferral_limit == 0 or len(grad_output_buffer) < wgrad_deferral_limit:
+                grad_output_buffer.append(grad_output)
+                wgrad_compute = False
+
+        if wgrad_compute:
+            if ctx.sequence_parallel:
+                world_size = get_tensor_model_parallel_world_size()
+                dim_size = list(input.size())
+                dim_size[0] = dim_size[0] * world_size
+
+                all_gather_buffer = get_global_memory_buffer().get_tensor(
+                    dim_size, input.dtype, "mpu"
+                )
+                handle = torch.distributed._all_gather_base(
+                    all_gather_buffer, input, group=get_tensor_model_parallel_group(), async_op=True
+                )
+                handle.wait()
+                # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
+                # gather is scheduled before the input gradient computation
+                total_input = all_gather_buffer
+            else:
+                total_input = input
+        grad_input = grad_output.matmul(weight)
+
+        # if ctx.sequence_parallel and wgrad_compute:
+        #     handle.wait()
+
+        if wgrad_compute:
+            grad_output, total_input = prepare_input_tensors_for_wgrad_compute(
+                grad_output, total_input
+            )
+
+        if ctx.allreduce_dgrad:
+            # Asynchronous all-reduce
+            handle = torch.distributed.all_reduce(
+                grad_input, group=get_tensor_model_parallel_group(), async_op=True
+            )
+            # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
+            # all-reduce is scheduled before the weight gradient computation
+
+        if ctx.sequence_parallel:
+            assert not ctx.allreduce_dgrad
+            dim_size = list(input.size())
+            sub_grad_input = torch.empty(
+                dim_size, dtype=input.dtype, device=torch.cuda.current_device(), requires_grad=False
+            )
+            # reduce_scatter
+            handle = torch.distributed._reduce_scatter_base(
+                sub_grad_input, grad_input, group=get_tensor_model_parallel_group(), async_op=True
+            )
+            handle.wait()
+            # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
+            # reduce scatter is scheduled before the weight gradient computation
+
+        if ctx.gradient_accumulation_fusion:
+            if wgrad_compute:
+                if weight.main_grad.dtype == torch.float32:
+                    fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32(
+                        total_input, grad_output, weight.main_grad
+                    )
+                elif weight.main_grad.dtype in (torch.float16, torch.bfloat16):
+                    fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp16(
+                        total_input, grad_output, weight.main_grad
+                    )
+                else:
+                    raise RuntimeError("Unsupported gradient type for gradient accumulation fusion")
+            # torch.addmm(weight.main_grad, grad_output.t(), total_input, out=weight.main_grad)
+            if hasattr(weight, 'grad_added_to_main_grad'):
+                # When overlap_grad_reduce is True, need to ensure that backward hooks
+                # are all run on the main backprop thread to prevent deadlocks. Setup
+                # dummy grad_weight tensor to prevent backward hooks from being run
+                # in a background thread.
+                if getattr(weight, 'zero_out_wgrad', False):
+                    grad_weight = torch.zeros(
+                        weight.main_grad.shape,
+                        dtype=input.dtype,
+                        device=torch.cuda.current_device(),
+                        requires_grad=False,
+                    )
+                else:
+                    grad_weight = torch.empty(
+                        weight.main_grad.shape,
+                        dtype=input.dtype,
+                        device=torch.cuda.current_device(),
+                        requires_grad=False,
+                    )
+                weight.grad_added_to_main_grad = True
+            else:
+                grad_weight = None
+        else:
+            grad_weight = grad_output.t().matmul(total_input)
+        grad_bias = grad_output.sum(dim=0) if use_bias else None
+
+        if ctx.sequence_parallel:
+            # handle.wait()
+            # Need to return None's as gradient has to flow for all the input arguments
+            # provided during forward
+            return sub_grad_input, grad_weight, grad_bias, None, None, None, None, None
+
+        if ctx.allreduce_dgrad:
+            handle.wait()
+        return grad_input, grad_weight, grad_bias, None, None, None, None, None
+
+
+import megatron.core.tensor_parallel.layers
+megatron.core.tensor_parallel.layers.LinearWithGradAccumulationAndAsyncCommunication = LinearWithGradAccumulationAndAsyncCommunication
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/moe_utils.py b/megatron-lm-musa-patch/musa_patch/moe_utils.py
new file mode 100644
index 000000000..02b2e7d48
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/moe_utils.py
@@ -0,0 +1,244 @@
+
+import math
+from typing import Optional, List
+
+import torch
+
+from megatron.core import parallel_state
+import megatron.core.transformer.moe.moe_utils
+from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region
+
+get_capacity = megatron.core.transformer.moe.moe_utils.get_capacity
+group_limited_topk = megatron.core.transformer.moe.moe_utils.group_limited_topk
+
+
+def sequence_load_balancing_loss_func(
+    probs: torch.Tensor,
+    routing_map: torch.Tensor,
+    batch_size: int,
+    seq_length: int,
+    topk: int,
+    moe_aux_loss_coeff: float,
+    sequence_partition_group=None,
+    moe_device_level_aux_loss_coeff: float=None,
+    moe_comm_aux_loss_coeff: float=None,
+    moe_router_topk_limited_devices: float=None,
+    moe_complementary_seq_aux_loss: bool=False,
+):
+    """
+    Calculate the auxiliary loss in sequence-level by computing the loss for each individual sample.
+    Refer to the DeepSeek-V2 huggingface repo
+    (https://huggingface.co/deepseek-ai/DeepSeek-V2) for details.
+
+    Args:
+        probs (torch.Tensor): Softmax probabilities output by the router for each token.
+                              Shape in [num_tokens, num_experts].
+        routing_map (torch.Tensor): Mapping of tokens to experts assignment.
+                                    Shape in [num_tokens, num_experts].
+        batch_size (int): Batch size to process.
+        seq_length (int): Sequence length to process.
+        topk (int): Number of experts to route to for each token.
+        moe_aux_loss_coeff (float): Scaling coefficient for the auxiliary loss.
+        sequence_partition_group (optional): The parallel group over which the sequence is
+                                             partitioned. If None, no partitioning is applied.
+                                             Defaults to None.
+
+    Returns:
+        torch.Tensor: The sequence auxiliary loss for load balancing.
+    """
+    num_sub_sequence = 1
+    num_experts = probs.shape[1]
+
+    probs_for_aux_loss = probs.view(seq_length, batch_size, -1)
+    routing_map = routing_map.view(seq_length, batch_size, -1)
+
+    # If the sequence is partitioned by certain parallelism strategies like Sequence Parallelism
+    # or Context Parallelism, compute the gradient of the auxiliary loss with respect to the full
+    # sequence.
+    if sequence_partition_group is not None:
+        num_sub_sequence = torch.distributed.get_world_size(sequence_partition_group)
+        seq_length *= num_sub_sequence
+        probs_for_aux_loss = gather_from_sequence_parallel_region(
+            probs_for_aux_loss, group=sequence_partition_group
+        )
+
+    cost_coeff = routing_map.sum(dim=0, dtype=torch.float).div_(seq_length * topk / num_experts)
+    if moe_complementary_seq_aux_loss:
+        assert (
+            (moe_device_level_aux_loss_coeff is None) and
+            (moe_comm_aux_loss_coeff is None)
+            ), "moe_complementary_seq_aux_loss only used in deepseekV3, which means no other aux loss used"
+        sum_value = probs_for_aux_loss.sum(dim=-1, keepdim=True)
+        probs_for_aux_loss = probs_for_aux_loss / (sum_value + 1e-20)
+    seq_aux_loss = (cost_coeff * probs_for_aux_loss.mean(dim=0)).sum(dim=1).mean()
+    seq_aux_loss *= moe_aux_loss_coeff
+
+    if moe_device_level_aux_loss_coeff is not None:
+        num_group = (
+        parallel_state.get_expert_model_parallel_world_size()
+        )  # num_group equals to expert parallel size
+        device_aux_loss = (cost_coeff.view(batch_size, num_group, -1).mean(dim=2) *
+                           probs_for_aux_loss.mean(dim=0).view(batch_size, num_group, -1).sum(dim=2)).sum(dim=1).mean()
+        device_aux_loss *= moe_device_level_aux_loss_coeff
+        seq_aux_loss += device_aux_loss
+    if moe_comm_aux_loss_coeff is not None:
+        num_group = (
+        parallel_state.get_expert_model_parallel_world_size()
+        )  # num_group equals to expert parallel size
+        cost_coeff = routing_map.view(seq_length, batch_size, num_group, -1).any(dim=3).sum(dim=0).float()
+        cost_coeff.div_(seq_length *  moe_router_topk_limited_devices / num_group)
+        comm_aux_loss = (cost_coeff *
+                           probs_for_aux_loss.mean(dim=0).view(batch_size, num_group, -1).sum(dim=2)).sum(dim=1).mean()
+        comm_aux_loss *= moe_comm_aux_loss_coeff
+        seq_aux_loss += comm_aux_loss
+
+    return seq_aux_loss
+
+def topk_softmax_with_capacity(
+    logits: torch.Tensor,
+    topk: int,
+    capacity_factor: Optional[float] = None,
+    pad_to_capacity: bool = False,
+    drop_policy: str = "probs",
+    use_pre_softmax: bool = False,
+    num_groups: Optional[int] = None,
+    group_topk: Optional[int] = None,
+    scaling_factor: Optional[float] = None,
+    deterministic_mode: bool = False,
+    score_function: str = "softmax",
+    expert_bias: Optional[torch.Tensor] = None,
+    device_level_capacity: bool = False,
+):
+    """Apply capacity and padding to the top-k selection.
+    Args:
+        logits (torch.Tensor): Logits tensor.
+        topk (int): The number of experts to select for each token.
+        capacity_factor (float): The capacity factor of each expert. Will drop tokens if the number
+                               of tokens exceeds the capacity.
+        pad_to_capacity (bool): Whether to need padding in token drop mode. The probs for padded
+                               tokens will be 0.
+        drop_policy (str): The policy to drop tokens. Can be either "prob" or "position".
+                           If "prob", the tokens with the lowest probabilities will be dropped.
+                           If "position", tokens at the end of each batch will be dropped.
+        use_pre_softmax (bool): Whether to apply softmax before top-k selection.
+        num_groups (int): Number of groups for routed experts.
+        group_topk (int): Number of selected groups for each token.
+        scaling_factor (float): Scaling factor of routing score in top-k selection.
+        deterministic_mode (bool): Deprecated.
+        score_function (str): The score function to use. Can be either "softmax" or "sigmoid".
+        expert_bias (torch.Tensor): The bias added to logits for expert routing.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            - routing_probs (torch.Tensor): A tensor of shape [num_tokens, num_experts] containing
+              the routing probabilities for each token to each expert.
+            - routing_map (torch.Tensor): A mask tensor of shape [num_tokens, num_experts]
+              indicating which experts were selected for each token. True values represent
+              the selected experts.
+            - tokens_per_expert (torch.Tensor): A tensor of shape [num_experts] containing
+              the number of local tokens assigned to each expert before dropping and padding.
+    """
+    assert logits.dim() == 2, f"Expected 2D logits [num_tokens, num_experts], got {logits.dim()}."
+    num_tokens, num_experts = logits.shape
+
+    def compute_topk(scores, topk, num_groups=None, group_topk=None):
+        if group_topk:
+            return group_limited_topk(
+                scores=scores,
+                topk=topk,
+                num_tokens=num_tokens,
+                num_experts=num_experts,
+                num_groups=num_groups,
+                group_topk=group_topk,
+            )
+        else:
+            return torch.topk(scores, k=topk, dim=1)
+
+    if score_function == "softmax":
+        if use_pre_softmax:
+            scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits)
+            probs, top_indices = compute_topk(scores, topk, num_groups, group_topk)
+        else:
+            scores, top_indices = compute_topk(logits, topk, num_groups, group_topk)
+            probs = torch.softmax(scores, dim=-1, dtype=torch.float32).type_as(logits)
+    elif score_function == "sigmoid":
+        scores = torch.sigmoid(logits)
+        if expert_bias is not None:
+            scores_for_routing = scores + expert_bias
+            _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk)
+            scores = torch.gather(scores, dim=1, index=top_indices).type_as(logits)
+        else:
+            scores, top_indices = compute_topk(scores, topk, num_groups, group_topk)
+        probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores
+    else:
+        raise ValueError(f"Invalid score_function: {score_function}")
+
+    if scaling_factor:
+        probs = probs * scaling_factor
+
+    # TODO Try using element-wise operations instead of scatter?
+    topk_masked_gates = torch.zeros_like(logits).scatter(1, top_indices, probs)
+    topk_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
+    tokens_per_expert = topk_map.sum(dim=0)
+
+    if capacity_factor is None:
+        # TopK without capacity
+        return topk_masked_gates, topk_map, tokens_per_expert
+    elif device_level_capacity:
+        assert drop_policy=='probs', f"only support 'probs' for device_level capacity, but get {drop_policy}"
+        num_group = (
+        parallel_state.get_expert_model_parallel_world_size()
+        )  # num_group equals to expert parallel size
+        device_expert_capacity = get_capacity(
+            num_tokens=num_tokens * topk, num_experts=num_experts, capacity_factor=capacity_factor
+        )*num_experts//num_group
+        # Maskout exceeded tokens
+        if drop_policy == "probs":
+            topk_masked_group_gates = topk_masked_gates.view(num_tokens, num_group, -1)
+            topk_masked_group_gates = topk_masked_group_gates.permute(0,2,1).reshape(-1, num_group)
+            _, capacity_indices = torch.topk(
+                topk_masked_group_gates, k=device_expert_capacity, dim=0, sorted=False
+            )
+            capacity_mask = torch.zeros([num_tokens*num_experts//num_group, num_group], device=logits.device).scatter(0, capacity_indices, 1).bool()
+            capacity_mask = capacity_mask.view(num_tokens, num_experts//num_group, num_group).permute(0,2,1).reshape(num_tokens, -1)
+        else:
+            raise ValueError(f"Invalid drop_policy: {drop_policy}")
+
+        if pad_to_capacity:
+            final_map = capacity_mask
+            final_probs = topk_masked_gates * final_map
+        else:
+            # Get exceed mask and maskout exceeded probs and indices
+            final_map = torch.logical_and(topk_map, capacity_mask)
+            final_probs = topk_masked_gates * final_map
+        return final_probs, final_map, tokens_per_expert
+    else:
+        # TopK with capacity
+        expert_capacity = get_capacity(
+            num_tokens=num_tokens * topk, num_experts=num_experts, capacity_factor=capacity_factor
+        )
+
+        # Maskout exceeded tokens
+        if drop_policy == "probs":
+            _, capacity_indices = torch.topk(
+                topk_masked_gates, k=expert_capacity, dim=0, sorted=False
+            )
+            capacity_mask = torch.zeros_like(logits).scatter(0, capacity_indices, 1).bool()
+        elif drop_policy == "position":
+            _, capacity_indices = torch.topk(topk_map.int(), k=expert_capacity, dim=0, sorted=False)
+            capacity_mask = torch.zeros_like(logits).scatter(0, capacity_indices, 1).bool()
+        else:
+            raise ValueError(f"Invalid drop_policy: {drop_policy}")
+
+        if pad_to_capacity:
+            final_map = capacity_mask
+            final_probs = topk_masked_gates * final_map
+        else:
+            # Get exceed mask and maskout exceeded probs and indices
+            final_map = torch.logical_and(topk_map, capacity_mask)
+            final_probs = topk_masked_gates * final_map
+        return final_probs, final_map, tokens_per_expert
+
+
+megatron.core.transformer.moe.moe_utils.sequence_load_balancing_loss_func = sequence_load_balancing_loss_func
+megatron.core.transformer.moe.moe_utils.topk_softmax_with_capacity = topk_softmax_with_capacity
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/optimizer.py b/megatron-lm-musa-patch/musa_patch/optimizer.py
new file mode 100644
index 000000000..7ac548516
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/optimizer.py
@@ -0,0 +1,238 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import logging
+import warnings
+from typing import Callable, Dict, List, Optional, Tuple
+
+import os
+import torch
+from torch.optim import SGD as CPUSGD
+
+try:
+    from transformer_engine.pytorch.optimizers import FusedAdam as Adam
+    from transformer_engine.pytorch.optimizers import FusedSGD as SGD
+except ImportError:
+    try:
+        from apex.optimizers import FusedAdam as Adam
+        from apex.optimizers import FusedSGD as SGD
+    except ImportError:
+        import warnings
+
+        warnings.warn(
+            f'Transformer Engine and Apex are not installed. Falling back to Torch optimizers.'
+        )
+
+        # Apex's FusedAdam is a drop-in replacement for torch's AdamW.
+        # pylint: disable-next=line-too-long.
+        # See https://github.com/NVIDIA/apex/blob/7b73b12361068a10b0f44844534613f252a5ea75/apex/optimizers/fused_adam.py#L16.
+        from torch.optim import AdamW as Adam, SGD
+
+from megatron.core.optimizer.cpu_offloading.hybrid_optimizer import HybridDeviceOptimizer
+from megatron.core.distributed.param_and_grad_buffer import _ParamAndGradBuffer
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer
+from megatron.core.optimizer.grad_scaler import ConstantGradScaler, DynamicGradScaler
+from megatron.core.optimizer import (
+    Float16OptimizerWithFloat16Params,
+    FP32Optimizer,
+    MegatronOptimizer,
+)
+from megatron.core.optimizer.optimizer_config import OptimizerConfig
+from megatron.core.utils import is_te_min_version
+
+
+logger = logging.getLogger(__name__)
+
+def _get_megatron_optimizer_based_on_param_groups(
+    config: OptimizerConfig,
+    model_chunks: List[MegatronModule],
+    param_groups: List,
+    per_model_buffers: Optional[Dict[int, List[_ParamAndGradBuffer]]] = None,
+    model_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
+    data_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
+    data_parallel_group_gloo: Optional[torch.distributed.ProcessGroup] = None,
+    data_parallel_group_idx: Optional[int] = None,
+    distributed_optimizer_instance_id: Optional[int] = 0,
+) -> MegatronOptimizer:
+    """Get Megatron optimizer based on parameter groups.
+
+    Args:
+        config (OptimizerConfig): optimizer configuration object.
+        model_chunks (list): list of model chunks.
+        param_groups (list): list of parameter groups.
+        per_model_buffers (dict, optional): buffers for distributed optimizer. Defaults to None.
+        data_parallel_group (torch.distributed.ProcessGroup, optional): data-parallel group for
+            distributed optimizer. Defaults to None.
+        data_parallel_group_gloo (torch.distributed.ProcessGroup, optional): gloo data-parallel
+            group for distributed optimizer. Defaults to None.
+        data_parallel_group_idx (int, optional): data-parallel group index for distributed
+            optimizer. Defaults to None.
+        distributed_optimizer_instance_id (int, optional): Distributed optimizer instance. Defaults
+            0.
+
+    Returns:
+        Instance of MegatronOptimizer.
+    """
+    # when freezing sub-models we may have no trainable parameters on a rank and
+    # hence an empty param_groups. However, we still need to create an optimizer
+    # for the purposes of grad stats reductions
+    if param_groups:
+        if config.optimizer_cpu_offload:
+            if torch.__version__ < '2.3.0':
+                # is_available = lambda: False
+                # set cuda not available for complex cuda inspection
+                torch.cuda.is_available = lambda : False
+                # use DeepSpeedCPUAdam for better performance
+                from deepspeed.ops.adam import DeepSpeedCPUAdam as CPUAdam
+                # reset the cuda availability back to normal
+                torch.cuda.is_available = torch.musa.is_available
+                warnings.warn("We use DeepSpeedCPUAdam instead of torch.optim.AdamW "
+                              "for better performace if torch.version < 2.3.0.")
+            else:
+                # torch.optim.AdamW supports __fused_adamw when torch.version >= 2.3.0
+                from torch.optim import AdamW as CPUAdam
+
+            # cpu optimizer offload must config use_precision_aware_optimizer to True,
+            # we should reconfig use_precision_aware_optimizer to break the compatibility.
+            if not int(os.getenv("CPU_OPTIMIZER_PRECISION_AWARE_RECONFIG", 0)):
+                config.use_precision_aware_optimizer = False
+
+            gpu_optimizer_cls = Adam if config.optimizer == 'adam' else SGD
+            cpu_optimizer_cls = CPUAdam if config.optimizer == 'adam' else CPUSGD
+            if config.use_torch_optimizer_for_cpu_offload:
+                gpu_optimizer_cls = cpu_optimizer_cls
+            if config.optimizer == 'adam':
+                gpu_optimizer_cls = Adam
+                cpu_optimizer_cls = CPUAdam
+                optimizer_defaults = dict(
+                    lr=config.lr,
+                    weight_decay=config.weight_decay,
+                    betas=(config.adam_beta1, config.adam_beta2),
+                    eps=config.adam_eps,
+                    bias_correction=True,
+                    fused=True,  # this flag is used to improve the performance of the cpu optimizer
+                )
+            else:
+                gpu_optimizer_cls = SGD
+                cpu_optimizer_cls = CPUSGD
+                optimizer_defaults = dict(
+                    lr=config.lr, weight_decay=config.weight_decay, momentum=config.sgd_momentum
+                )
+
+            optimizer = HybridDeviceOptimizer(
+                param_groups,
+                offload_fraction=config.optimizer_offload_fraction,
+                cpu_optimizer_cls=cpu_optimizer_cls,
+                gpu_optimizer_cls=gpu_optimizer_cls,
+                overlap_cpu_optimizer_d2h_h2d=config.overlap_cpu_optimizer_d2h_h2d,
+                pin_cpu_grads=config.pin_cpu_grads,
+                pin_cpu_params=config.pin_cpu_params,
+                param_update_in_fp32=True,
+                **optimizer_defaults,
+            )
+            init_state_fn = None
+        elif config.optimizer == 'adam':
+            kwargs = {
+                "params": param_groups,
+                "lr": config.lr,
+                "weight_decay": config.weight_decay,
+                "betas": (config.adam_beta1, config.adam_beta2),
+                "eps": config.adam_eps,
+            }
+
+            if config.use_precision_aware_optimizer:
+                kwargs.update(
+                    {
+                        "master_weights": True,
+                        "use_decoupled_grad": True,
+                        "master_weight_dtype": config.main_params_dtype,
+                        "exp_avg_dtype": config.exp_avg_dtype,
+                        "exp_avg_sq_dtype": config.exp_avg_sq_dtype,
+                    }
+                )
+
+                if is_te_min_version("2.1.0.dev0"):
+                    kwargs.update({"store_param_remainders": True})
+
+            optimizer = Adam(**kwargs)
+
+            def init_state_fn(opt, config=None):
+                for group in opt.param_groups:
+                    for p in group['params']:
+                        if len(opt.state[p]) == 0:
+                            if config is None or not config.use_precision_aware_optimizer:
+                                opt.state[p]['exp_avg'] = torch.zeros_like(p.data)
+                                opt.state[p]['exp_avg_sq'] = torch.zeros_like(p.data)
+                            else:
+                                opt.initialize_state(p)
+
+        elif config.optimizer == 'sgd':
+            optimizer = SGD(
+                param_groups,
+                lr=config.lr,
+                weight_decay=config.weight_decay,
+                momentum=config.sgd_momentum,
+            )
+            init_state_fn = None
+        else:
+            raise Exception('{} optimizer is not supported.'.format(config.optimizer))
+    else:
+        optimizer = None
+        init_state_fn = None
+
+    # Mixed precision optimizer.
+    # - Note: both the Float16Optimizer and the DistributedOptimizer inherit
+    #   from the MixedPrecisionOptimizer, which manages any optimizer where
+    #   the model params and main params are distinct.
+    if config.fp16 or config.bf16 or config.use_distributed_optimizer:
+
+        # Grad scaler:
+        #    if loss-scale is provided, instantiate the constant scaler.
+        #    if we are using fp16 and loss-scale is not present, use a
+        #       dynamic scaler.
+        #    otherwise we are running in bf16 with no loss-scale so
+        #       leave it as None.
+        grad_scaler = None
+
+        # Constant loss scale.
+        if config.loss_scale:
+            grad_scaler = ConstantGradScaler(config.loss_scale)
+
+        # Dynamic loss scale.
+        else:
+            if config.fp16:
+                grad_scaler = DynamicGradScaler(
+                    initial_scale=config.initial_loss_scale,
+                    min_scale=config.min_loss_scale,
+                    growth_factor=2.0,
+                    backoff_factor=0.5,
+                    growth_interval=config.loss_scale_window,
+                    hysteresis=config.hysteresis,
+                )
+
+        optimizer_args = [optimizer, config, grad_scaler, init_state_fn]
+        if config.use_distributed_optimizer:
+            optimizer = DistributedOptimizer(
+                *optimizer_args,
+                model_chunks=model_chunks,
+                per_model_buffers=per_model_buffers,
+                data_parallel_group=data_parallel_group,
+                data_parallel_group_gloo=data_parallel_group_gloo,
+                data_parallel_group_idx=data_parallel_group_idx,
+                distributed_optimizer_instance_id=distributed_optimizer_instance_id,
+            )
+        else:
+            optimizer = Float16OptimizerWithFloat16Params(*optimizer_args)
+            setattr(optimizer, 'grad_stats_parallel_group', model_parallel_group)
+    else:
+        # FP32 optimizer.
+        optimizer = FP32Optimizer(optimizer, config, init_state_fn)
+        setattr(optimizer, 'grad_stats_parallel_group', model_parallel_group)
+
+    return optimizer
+
+import sys
+for k in sys.modules:
+    if k.startswith('megatron'):
+        for target in ['_get_megatron_optimizer_based_on_param_groups']:
+            if getattr(sys.modules[k], target, None):
+                setattr(sys.modules[k], target, _get_megatron_optimizer_based_on_param_groups)
diff --git a/megatron-lm-musa-patch/musa_patch/p2p_communication.py b/megatron-lm-musa-patch/musa_patch/p2p_communication.py
new file mode 100644
index 000000000..89ad8c12e
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/p2p_communication.py
@@ -0,0 +1,93 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+import operator
+from functools import reduce
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+
+from megatron import core
+from megatron.core import ModelParallelConfig
+from megatron.core.parallel_state import (
+    get_pipeline_model_parallel_group,
+    get_pipeline_model_parallel_next_rank,
+    get_pipeline_model_parallel_prev_rank,
+    get_pipeline_model_parallel_rank,
+    get_pipeline_model_parallel_world_size,
+)
+
+def _p2p_ops(
+    *,
+    tensor_send_prev: Optional[torch.Tensor],
+    tensor_recv_prev: Optional[torch.Tensor],
+    tensor_send_next: Optional[torch.Tensor],
+    tensor_recv_next: Optional[torch.Tensor],
+    group: torch.distributed.ProcessGroup,
+    prev_pipeline_rank: int,
+    next_pipeline_rank: int,
+):
+    reqs = []
+    rank = get_pipeline_model_parallel_rank()
+    even_send_odd_recv_group = group
+    # if get_pipeline_model_parallel_world_size() == 2:
+    #     # Use the global process group for one of the two p2p communications
+    #     # to allow the overlap of the independent communications.
+    #     # Using the global process group is compatible because the pipeline-parallel
+    #     # communications set the source and destination by global rank.
+    #     even_recv_odd_send_group = torch.distributed.group.WORLD
+    # else:
+    even_recv_odd_send_group = group
+
+    if get_pipeline_model_parallel_rank() % 2 == 0:
+        if tensor_send_next is not None:
+            send_next_req = torch.distributed.isend(
+                tensor=tensor_send_next, dst=next_pipeline_rank, group=even_send_odd_recv_group
+            )
+            reqs.append(send_next_req)
+
+        if tensor_recv_prev is not None:
+            recv_prev_req = torch.distributed.irecv(
+                tensor=tensor_recv_prev, src=prev_pipeline_rank, group=even_recv_odd_send_group
+            )
+            reqs.append(recv_prev_req)
+
+        if tensor_send_prev is not None:
+            send_prev_req = torch.distributed.isend(
+                tensor=tensor_send_prev, dst=prev_pipeline_rank, group=even_send_odd_recv_group
+            )
+            reqs.append(send_prev_req)
+
+        if tensor_recv_next is not None:
+            recv_next_req = torch.distributed.irecv(
+                tensor=tensor_recv_next, src=next_pipeline_rank, group=even_recv_odd_send_group
+            )
+            reqs.append(recv_next_req)
+
+    else:
+        if tensor_recv_prev is not None:
+            recv_prev_req = torch.distributed.irecv(
+                tensor=tensor_recv_prev, src=prev_pipeline_rank, group=even_send_odd_recv_group
+            )
+            reqs.append(recv_prev_req)
+
+        if tensor_send_next is not None:
+            send_next_req = torch.distributed.isend(
+                tensor=tensor_send_next, dst=next_pipeline_rank, group=even_recv_odd_send_group
+            )
+            reqs.append(send_next_req)
+
+        if tensor_recv_next is not None:
+            recv_next_req = torch.distributed.irecv(
+                tensor=tensor_recv_next, src=next_pipeline_rank, group=even_send_odd_recv_group
+            )
+            reqs.append(recv_next_req)
+
+        if tensor_send_prev is not None:
+            send_prev_req = torch.distributed.isend(
+                tensor=tensor_send_prev, dst=prev_pipeline_rank, group=even_recv_odd_send_group
+            )
+            reqs.append(send_prev_req)
+    return reqs
+
+import megatron.core.pipeline_parallel.p2p_communication
+megatron.core.pipeline_parallel.p2p_communication._p2p_ops = _p2p_ops
diff --git a/megatron-lm-musa-patch/musa_patch/parallel_state.py b/megatron-lm-musa-patch/musa_patch/parallel_state.py
new file mode 100644
index 000000000..ed163edec
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/parallel_state.py
@@ -0,0 +1,823 @@
+import os
+import sys
+import logging
+import warnings
+from datetime import timedelta
+from functools import partial
+from itertools import cycle
+from typing import Callable, List, Optional
+
+import torch
+from megatron.core.parallel_state import *
+import megatron.core.parallel_state as parallel_state
+
+logger = logging.getLogger(__name__)
+
+_EPX_DATA_PARALLEL_LCP = None
+
+globals().update({k: getattr(parallel_state, k) for k in dir(parallel_state) if k.startswith('_')})
+
+group_list = {
+    name: value for name, value in globals().items()
+    if name.startswith("_") and not callable(value)
+}
+
+def initialize_model_parallel(
+    tensor_model_parallel_size: int = 1,
+    pipeline_model_parallel_size: int = 1,
+    virtual_pipeline_model_parallel_size: Optional[int] = None,
+    pipeline_model_parallel_split_rank: Optional[int] = None,
+    pipeline_model_parallel_comm_backend: Optional[str] = None,
+    use_sharp: bool = False,
+    context_parallel_size: int = 1,
+    hierarchical_context_parallel_sizes: Optional[List[int]] = None,
+    expert_model_parallel_size: int = 1,
+    num_distributed_optimizer_instances: int = 1,
+    expert_tensor_parallel_size: Optional[int] = None,
+    nccl_communicator_config_path: Optional[str] = None,
+    distributed_timeout_minutes: int = 30,
+    order: str = "tp-cp-ep-dp-pp",
+    encoder_tensor_model_parallel_size: int = 0,
+    encoder_pipeline_model_parallel_size: Optional[int] = 0,
+    get_embedding_ranks: Optional[Callable[[List[int], Optional[int]], List[int]]] = None,
+    get_position_embedding_ranks: Optional[Callable[[List[int], Optional[int]], List[int]]] = None,
+    create_gloo_process_groups: bool = True,
+) -> None:
+    # pylint: disable=line-too-long
+    """Initialize model data parallel groups.
+
+    Args:
+        tensor_model_parallel_size (int, default = 1):
+            The number of GPUs to split individual tensors across.
+
+        pipeline_model_parallel_size (int, default = 1):
+            The number of tensor parallel GPU groups to split the
+            Transformer layers across. For example, if
+            tensor_model_parallel_size is 4 and
+            pipeline_model_parallel_size is 2, the model will be split
+            into 2 groups of 4 GPUs.
+
+        virtual_pipeline_model_parallel_size (int, optional):
+            The number of stages that each pipeline group will have,
+            interleaving as necessary. If None, no interleaving is
+            performed. For example, if tensor_model_parallel_size is 1,
+            pipeline_model_parallel_size is 4,
+            virtual_pipeline_model_parallel_size is 2, and there are
+            16 transformer layers in the model, the model will be
+            split into 8 stages with two layers each and each GPU
+            would get 2 stages as such (layer number starting with 1):
+
+            GPU 0: [1, 2] [9, 10]
+            GPU 1: [3, 4] [11, 12]
+            GPU 2: [5, 6] [13, 14]
+            GPU 3: [7, 8] [15, 16]
+
+        pipeline_model_parallel_split_rank (int, optional):
+            DEPRECATED. For models with both an encoder and decoder, the rank in
+            pipeline to switch between encoder and decoder (i.e. the
+            first rank of the decoder). This allows the user to set
+            the pipeline parallel size of the encoder and decoder
+            independently. For example, if
+            pipeline_model_parallel_size is 8 and
+            pipeline_model_parallel_split_rank is 3, then ranks 0-2
+            will be the encoder and ranks 3-7 will be the decoder.
+
+        pipeline_model_parallel_comm_backend (str, optional):
+            The backend to use for pipeline parallel communication.
+            If None, the default backend will be used.
+
+        use_sharp (bool, default = False):
+            Set the use of SHARP for the collective communications of
+            data-parallel process groups. When `True`, run barrier
+            within each data-parallel process group, which specifies
+            the SHARP application target groups.
+
+        context_parallel_size (int, default = 1):
+            The number of tensor parallel GPU groups to split the
+            network input sequence length across. Compute of attention
+            module requires tokens of full sequence length, so GPUs
+            in a context parallel group need to communicate with each
+            other to exchange information of other sequence chunks.
+            Each GPU and its counterparts in other tensor parallel
+            groups compose a context parallel group.
+
+            For example, assume we have 8 GPUs, if tensor model parallel
+            size is 4 and context parallel size is 2, the network input
+            will be split into two sequence chunks, which are processed
+            by 2 different groups of 4 GPUs. One chunk is processed by
+            GPU0-3, the other chunk is processed by GPU4-7. Four groups
+            are build to do context parallel communications: [GPU0, GPU4],
+            [GPU1, GPU5], [GPU2, GPU6], and [GPU3, GPU7].
+
+            Context parallelism partitions sequence length, so it has no
+            impact on weights, which means weights are duplicated among
+            GPUs in a context parallel group. Hence, weight gradients
+            all-reduce is required in backward. For simplicity, we piggyback
+            GPUs of context parallelism on data parallel group for
+            weight gradient all-reduce.
+
+        expert_model_parallel_size (int, default = 1):
+            The number of Mixture of Experts parallel GPUs in each expert
+            parallel group.
+
+        num_distributed_optimizer_instances (int, default = 1):
+            The number of distributed optimizer replicas across the data-
+            parallel domain.
+
+        expert_tensor_parallel_size (int, default = tp_size):
+            The number of GPUs to split individual tensors of expert.
+
+        nccl_communicator_config_path (str, default = None):
+            Path to the yaml file of NCCL communicator configurations.
+            `min_ctas`, `max_ctas`, and `cga_cluster_size` can be set
+            for each communicator.
+
+        distributed_timeout_minutes (int, default = 30): Timeout, in
+            minutes,for operations executed against distributed
+            process groups. See PyTorch documentation at
+            https://pytorch.org/docs/stable/distributed.html for
+            caveats.
+
+        order (str, default=tp-dp-pp):
+            The rank initialization order of parallelism. Now we support
+            tp-dp-pp and tp-pp-dp orders.
+
+        encoder_tensor_model_parallel_size (int, default = 0):
+            The number of GPUs to split individual tensors across in the encoder. If 0,
+            then we use the default, decoder's tensor model parallel size.
+
+        encoder_pipeline_model_parallel_size (int, default = 0):
+            The number of tensor parallel GPU groups to allocate to the encoder. As an example,
+            if pipeline_model_parallel_size is 4 and encoder_pipeline_model_parallel_size is 2,
+            then the encoder will use the first two pipeline stages for its layers, and the total
+            amount of pipelineing is 6.
+
+        get_embedding_ranks (Callable[[List[int], Optional[int]], List[int]], optional, default=None):
+            A function that takes in a list of ranks for a pipeline group and returns
+            those ranks that should have embeddings.
+
+        get_position_embedding_ranks (Callable[[List[int], Optional[int]], List[int]], optional, default=None):
+            A function that takes in a list of ranks for a pipeline group, and returns
+            those ranks that should have position embeddings.
+
+        create_gloo_process_groups (bool, default = True):
+            Create Gloo process groups if set to True. If set to False, Gloo process groups are
+            not created and calls to get Gloo process groups will result in assertion errors.
+
+    Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
+    use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
+    the model pipeline. The present function will
+    create 8 tensor model-parallel groups, 4 pipeline model-parallel groups
+    and 8 data-parallel groups as:
+        8 data_parallel groups:
+            [g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15]
+        8 tensor model-parallel groups:
+            [g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15]
+        4 pipeline model-parallel groups:
+            [g0, g4, g8, g12], [g1, g5, g9, g13], [g2, g6, g10, g14], [g3, g7, g11, g15]
+    Note that for efficiency, the caller should make sure adjacent ranks
+    are on the same DGX box. For example if we are using 2 DGX-1 boxes
+    with a total of 16 GPUs, rank 0 to 7 belong to the first box and
+    ranks 8 to 15 belong to the second box.
+
+    """
+
+    if encoder_pipeline_model_parallel_size is None:
+        encoder_pipeline_model_parallel_size = 0
+
+    if encoder_tensor_model_parallel_size == 0 and encoder_pipeline_model_parallel_size > 0:
+        encoder_tensor_model_parallel_size = tensor_model_parallel_size
+
+    if get_embedding_ranks is None:
+        get_embedding_ranks = partial(
+            default_embedding_ranks, split_rank=pipeline_model_parallel_split_rank
+        )
+
+    if get_position_embedding_ranks is None:
+        get_position_embedding_ranks = partial(
+            default_position_embedding_ranks, split_rank=pipeline_model_parallel_split_rank
+        )
+
+    if encoder_pipeline_model_parallel_size > 0:
+        global _PIPELINE_MODEL_PARALLEL_DECODER_START
+        _PIPELINE_MODEL_PARALLEL_DECODER_START = encoder_pipeline_model_parallel_size
+
+    # Get world size and rank. Ensure some consistencies.
+    assert torch.distributed.is_initialized()
+    world_size: int = torch.distributed.get_world_size()
+
+    if encoder_tensor_model_parallel_size > 0:
+        assert (
+            encoder_tensor_model_parallel_size <= tensor_model_parallel_size
+        ), "We do not support encoders with more TP than the decoder."
+
+    encoder_model_size = (
+        encoder_tensor_model_parallel_size
+        * encoder_pipeline_model_parallel_size
+        * context_parallel_size
+    )
+    decoder_model_size = (
+        tensor_model_parallel_size * pipeline_model_parallel_size * context_parallel_size
+    )
+    total_model_size = encoder_model_size + decoder_model_size
+
+    if world_size % total_model_size != 0:
+        raise RuntimeError(f"world_size ({world_size}) is not divisible by {total_model_size}")
+
+    data_parallel_size: int = world_size // total_model_size
+
+    encoder_world_size = encoder_model_size * data_parallel_size
+    decoder_world_size = decoder_model_size * data_parallel_size
+
+    assert (
+        encoder_world_size + decoder_world_size == world_size
+    ), f"{encoder_world_size=} + {decoder_world_size=} != {world_size=}"
+
+    if virtual_pipeline_model_parallel_size is not None:
+        if not pipeline_model_parallel_size > 1:
+            raise RuntimeError(
+                "pipeline-model-parallel size should be greater than 1 with interleaved schedule"
+            )
+        global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
+        global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
+        _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = 0
+        _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = virtual_pipeline_model_parallel_size
+
+    if pipeline_model_parallel_split_rank is not None:
+        global _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
+        _PIPELINE_MODEL_PARALLEL_SPLIT_RANK = pipeline_model_parallel_split_rank
+
+    rank = torch.distributed.get_rank()
+
+    nccl_comm_cfgs = {}
+    if nccl_communicator_config_path is not None:
+        try:
+            import yaml
+        except ImportError:
+            raise RuntimeError(
+                "Cannot import `yaml`. Setting custom nccl communicator configs "
+                "requires the yaml package."
+            )
+
+        with open(nccl_communicator_config_path, "r") as stream:
+            nccl_comm_cfgs = yaml.safe_load(stream)
+
+    if encoder_world_size > 0:
+        encoder_rank_generator = RankGenerator(
+            tp=encoder_tensor_model_parallel_size,
+            ep=1,
+            dp=data_parallel_size,
+            pp=encoder_pipeline_model_parallel_size,
+            cp=context_parallel_size,
+            order=order,
+            rank_offset=0,
+        )
+    else:
+        encoder_rank_generator = None
+
+    decoder_rank_generator = RankGenerator(
+        tp=tensor_model_parallel_size,
+        ep=1,
+        dp=data_parallel_size,
+        pp=pipeline_model_parallel_size,
+        cp=context_parallel_size,
+        order=order,
+        rank_offset=encoder_world_size,
+    )
+
+    # Build expert rank generator
+    if expert_tensor_parallel_size is None:
+        expert_tensor_parallel_size = tensor_model_parallel_size
+    expert_tensor_model_pipeline_parallel_size = (
+        expert_tensor_parallel_size * expert_model_parallel_size * pipeline_model_parallel_size
+    )
+    expert_data_parallel_size = decoder_world_size // expert_tensor_model_pipeline_parallel_size
+    if decoder_world_size % expert_tensor_model_pipeline_parallel_size != 0:
+        raise RuntimeError(
+            f"decoder world_size ({decoder_world_size}) is not divisible by expert_tensor_model_pipeline_parallel size ({expert_tensor_model_pipeline_parallel_size})"
+        )
+
+    # TODO: support expert specific ordering
+    expert_decoder_rank_generator = RankGenerator(
+        tp=expert_tensor_parallel_size,
+        ep=expert_model_parallel_size,
+        dp=expert_data_parallel_size,
+        pp=pipeline_model_parallel_size,
+        cp=1,
+        order=order,
+        rank_offset=encoder_world_size,
+    )
+
+    assert (
+        order.endswith("pp")
+        or pipeline_model_parallel_size == 1
+        or expert_data_parallel_size == data_parallel_size
+    ), "When not using pp-last rank ordering, the data parallel size of the attention and moe layers must be the same"
+
+    assert decoder_rank_generator.get_ranks("pp") == expert_decoder_rank_generator.get_ranks(
+        "pp"
+    ), f"Pipeline parallel groups are expected to be the same for Non-Expert and Expert part, \
+    but got {decoder_rank_generator.get_ranks('pp')} and {expert_decoder_rank_generator.get_ranks('pp')}"
+
+    def generator_wrapper(group_type, is_expert=False, **kwargs):
+        """The `RankGenerator` class produces a hyper-rectangle for a given set of
+        tensor, pipeline, data, expert, and context parallelism. If we have an encoder,
+        in addition to the default decoder, we essentially instantiate two `RankGenerator`
+        classes to construct the parallelism for each module separately, and we then have
+        to stitch them together for the right groups. For now, this means pp and tp-pp."""
+        if is_expert:
+            d_ranks = expert_decoder_rank_generator.get_ranks(group_type, **kwargs)
+        else:
+            d_ranks = decoder_rank_generator.get_ranks(group_type, **kwargs)
+
+        if encoder_rank_generator is None:
+            for x in d_ranks:
+                yield x
+            return
+        e_ranks = encoder_rank_generator.get_ranks(group_type, **kwargs)
+        if group_type == 'pp':
+            # Map 1 encoder tp rank to several decoder tp ranks, because
+            # these won't be the same size.
+            for x, y in zip(cycle(e_ranks), d_ranks):
+                yield x + y
+        elif group_type == 'tp-pp':
+            # For this group, we can just return the concatenated
+            # groups together, because their sizes are the same.
+            assert len(e_ranks) == len(d_ranks)
+            for x, y in zip(e_ranks, d_ranks):
+                yield x + y
+        else:
+            for x in e_ranks:
+                yield x
+            for x in d_ranks:
+                yield x
+
+    timeout = timedelta(minutes=distributed_timeout_minutes)
+
+    # Build the data-parallel groups.
+    global _DATA_PARALLEL_GROUP
+    global _DATA_PARALLEL_GROUP_GLOO
+    global _DATA_PARALLEL_GLOBAL_RANKS
+    global _DATA_PARALLEL_GROUP_WITH_CP
+    global _DATA_PARALLEL_GROUP_WITH_CP_GLOO
+    global _DATA_PARALLEL_GLOBAL_RANKS_WITH_CP
+    global _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP
+    global _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP_GLOO
+    global _INTER_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP
+    assert _DATA_PARALLEL_GROUP is None, 'data parallel group is already initialized'
+
+    global _EPX_DATA_PARALLEL_LCP
+    if int(os.getenv("USE_EPX", 0)):
+        from epx.process_group import EpxProcessGroup
+        from epx.lcp import Lcp
+        import torch.distributed as dist
+
+        logger.info(f"start initialization _EPX_DATA_PARALLEL_LCP for epx")
+
+        epx_rank = int(os.environ.get("RANK", 0))
+
+        pg = EpxProcessGroup(group_name=str(epx_rank))
+
+        # rank = torch.distributed.get_rank()
+        _EPX_DATA_PARALLEL_LCP = Lcp(pg, rank)
+
+        logger.info(f"finish initialization _EPX_DATA_PARALLEL_LCP for epx")
+
+
+    for ranks in generator_wrapper('dp'):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('dp', nccl_comm_cfgs),
+            group_desc='DATA_PARALLEL_GROUP',
+        )
+        if create_gloo_process_groups:
+            group_gloo = create_group(
+                ranks, timeout=timeout, backend="gloo", group_desc='DATA_PARALLEL_GROUP_GLOO'
+            )
+        else:
+            group_gloo = None
+        if rank in ranks:
+            _DATA_PARALLEL_GROUP = group
+            _DATA_PARALLEL_GROUP_GLOO = group_gloo
+            _DATA_PARALLEL_GLOBAL_RANKS = ranks
+
+    assert (
+        data_parallel_size * context_parallel_size
+    ) % num_distributed_optimizer_instances == 0, (
+        'Data parallel size should be divisible by partial DistOpt shard factor'
+    )
+    intra_partial_data_parallel_size = (
+        data_parallel_size * context_parallel_size
+    ) // num_distributed_optimizer_instances
+
+    for ranks_with_cp in generator_wrapper('dp-cp'):
+        group_with_cp = create_group(
+            ranks_with_cp,
+            timeout=timeout,
+            pg_options=get_nccl_options('dp_cp', nccl_comm_cfgs),
+            group_desc='DATA_PARALLEL_GROUP_WITH_CP',
+        )
+        if create_gloo_process_groups:
+            group_with_cp_gloo = create_group(
+                ranks_with_cp,
+                timeout=timeout,
+                backend="gloo",
+                group_desc='DATA_PARALLEL_GROUP_WITH_CP_GLOO',
+            )
+        else:
+            group_with_cp_gloo = None
+        if rank in ranks_with_cp:
+            _DATA_PARALLEL_GROUP_WITH_CP = group_with_cp
+            _DATA_PARALLEL_GROUP_WITH_CP_GLOO = group_with_cp_gloo
+            _DATA_PARALLEL_GLOBAL_RANKS_WITH_CP = ranks_with_cp
+
+        if num_distributed_optimizer_instances > 1:
+            # Create groups for Partial DistOpt, one for intra-partial DP domain
+            # Another for inter-partial DP domain
+            for i in range(num_distributed_optimizer_instances):
+                intra_partial_data_parallel_ranks_with_cp = ranks_with_cp[
+                    (i * intra_partial_data_parallel_size) : (
+                        (i + 1) * intra_partial_data_parallel_size
+                    )
+                ]
+
+                intra_partial_data_parallel_group_with_cp = create_group(
+                    intra_partial_data_parallel_ranks_with_cp,
+                    timeout=timeout,
+                    pg_options=get_nccl_options('intra_dp_cp', nccl_comm_cfgs),
+                    group_desc='INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP',
+                )
+                if create_gloo_process_groups:
+                    intra_partial_data_parallel_group_with_cp_gloo = create_group(
+                        intra_partial_data_parallel_ranks_with_cp,
+                        timeout=timeout,
+                        backend="gloo",
+                        group_desc='INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP_GLOO',
+                    )
+                else:
+                    intra_partial_data_parallel_group_with_cp_gloo = None
+
+                if rank in intra_partial_data_parallel_ranks_with_cp:
+                    _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP = (
+                        intra_partial_data_parallel_group_with_cp
+                    )
+                    _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP_GLOO = (
+                        intra_partial_data_parallel_group_with_cp_gloo
+                    )
+
+            for i in range(intra_partial_data_parallel_size):
+                inter_partial_data_parallel_ranks_with_cp = ranks_with_cp[
+                    i::intra_partial_data_parallel_size
+                ]
+
+                inter_partial_data_parallel_group_with_cp = create_group(
+                    inter_partial_data_parallel_ranks_with_cp,
+                    timeout=timeout,
+                    pg_options=get_nccl_options('inter_dp_cp', nccl_comm_cfgs),
+                    group_desc='INTER_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP',
+                )
+
+                if rank in inter_partial_data_parallel_ranks_with_cp:
+                    _INTER_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP = (
+                        inter_partial_data_parallel_group_with_cp
+                    )
+        else:
+            _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP = _DATA_PARALLEL_GROUP_WITH_CP
+            _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP_GLOO = _DATA_PARALLEL_GROUP_WITH_CP_GLOO
+
+    # Apply SHARP to DP process groups
+    if use_sharp:
+        if rank == 0:
+            print(
+                "The number of process groups to use SHARP with depends on the type "
+                "of the network switch. Nvidia QM1 switch supports SAHRP up to 8 "
+                "process groups and QM2 supports up to 256 process groups. We apply "
+                "SHARP to the communications of the data-parallel domain. If the "
+                "number of data-parallel process groups is larger than the max "
+                "process groups that the network switch supports, the communication "
+                "will fall back to non-SHARP operators. To enable SHARP, "
+                "`#SBATCH_NETWORK=sharp` should be set in the sbatch script."
+            )
+        torch.distributed.barrier(
+            group=get_data_parallel_group(with_context_parallel=True),
+            device_ids=[torch.cuda.current_device()],
+        )
+        # Set `NCCL_COLLNET_ENABLE=0` to restrict SHARP application to DP process groups
+        os.environ["NCCL_COLLNET_ENABLE"] = "0"
+
+    # Build the context-parallel groups.
+    global _CONTEXT_PARALLEL_GROUP
+    global _CONTEXT_PARALLEL_GLOBAL_RANKS
+    assert _CONTEXT_PARALLEL_GROUP is None, 'context parallel group is already initialized'
+    for ranks in generator_wrapper('cp'):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('cp', nccl_comm_cfgs),
+            group_desc='CONTEXT_PARALLEL_GROUP',
+        )
+        if rank in ranks:
+            _CONTEXT_PARALLEL_GROUP = group
+            _CONTEXT_PARALLEL_GLOBAL_RANKS = ranks
+        if hierarchical_context_parallel_sizes:
+            global _HIERARCHICAL_CONTEXT_PARALLEL_GROUPS
+            _HIERARCHICAL_CONTEXT_PARALLEL_GROUPS += create_hierarchical_parallel_groups(
+                rank,
+                ranks,
+                context_parallel_size,
+                hierarchical_context_parallel_sizes,
+                get_nccl_options('hcp', nccl_comm_cfgs),
+            )
+
+    # Build the model-parallel groups.
+    global _MODEL_PARALLEL_GROUP
+    global _MODEL_PARALLEL_GLOBAL_RANKS
+    assert _MODEL_PARALLEL_GROUP is None, 'model parallel group is already initialized'
+    for ranks in generator_wrapper('tp-pp'):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('mp', nccl_comm_cfgs),
+            group_desc='MODEL_PARALLEL_GROUP',
+        )
+        if rank in ranks:
+            _MODEL_PARALLEL_GROUP = group
+            _MODEL_PARALLEL_GLOBAL_RANKS = ranks
+
+    # Build the tensor model-parallel groups.
+    global _TENSOR_MODEL_PARALLEL_GROUP
+    global _TENSOR_MODEL_PARALLEL_GLOBAL_RANKS
+    assert (
+        _TENSOR_MODEL_PARALLEL_GROUP is None
+    ), 'tensor model parallel group is already initialized'
+    for ranks in generator_wrapper('tp'):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('tp', nccl_comm_cfgs),
+            group_desc='TENSOR_MODEL_PARALLEL_GROUP',
+        )
+        if rank in ranks:
+            _TENSOR_MODEL_PARALLEL_GROUP = group
+            _TENSOR_MODEL_PARALLEL_GLOBAL_RANKS = ranks
+
+    # Build the pipeline model-parallel groups and embedding groups
+    # (first and last rank in each pipeline model-parallel group).
+    global _PIPELINE_MODEL_PARALLEL_GROUP
+    global _PIPELINE_GLOBAL_RANKS
+    assert (
+        _PIPELINE_MODEL_PARALLEL_GROUP is None
+    ), 'pipeline model parallel group is already initialized'
+    global _EMBEDDING_GROUP
+    global _EMBEDDING_GLOBAL_RANKS
+    assert _EMBEDDING_GROUP is None, 'embedding group is already initialized'
+    global _POSITION_EMBEDDING_GROUP
+    global _POSITION_EMBEDDING_GLOBAL_RANKS
+    assert _POSITION_EMBEDDING_GROUP is None, 'position embedding group is already initialized'
+    if pipeline_model_parallel_comm_backend == 'ucc':
+        # The UCC backend provides two key benefits:
+        # 1) Achieves better bandwidth utilization than NCCL when using InfiniBand links.
+        # 2) Does not use GPU SM resources (Zero-SM), mitigating performance interference
+        #    with overlapping compute kernels.
+
+        # The UCC backend is recommended in the following cases:
+        # 1) When the exposed pipeline-parallel (PP) communications are significant.
+        #    - E.g., Pipeline parallelism with very less gradient accumulation steps.
+        #    - It may provide better performance due to improved bandwidth utilization.
+        # 2) When the critical-path pipeline stage has substantial PP-communication overlap.
+        #    - E.g., Uneven pipeline parallelism.
+        #    - It may provide better performance due to zero SM resource usage.
+        if 'CUDA_DEVICE_MAX_CONNECTIONS' in os.environ:
+            # UCC backend requires CUDA_DEVICE_MAX_CONNECTIONS variable to be larger than 1,
+            # to gurantee the overlapped UCC communications. If this environment variable is set to 1,
+            # all the UCC communication will be serialized.
+            assert (
+                os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] != '1'
+            ), "UCC-backend requires CUDA_DEVICE_MAX_CONNECTIONS > 1"
+
+        # Setting up required environment variables for ucc backend
+        #
+        # "TORCH_UCC_BLOCKING_WAIT=none" allows non-blocking waits of the communiction handle
+        # "UCC_EC_CUDA_STREAM_TASK_MODE" controls how CUDA execution engines (EC)
+        # schedule tasks on CUDA streams.
+        # "UCX_TLS" controls transport layer selection
+        # "NSYS_UCP_COMM_PARAMS=1" enables capturing ucx tracing in nsys profiling
+        # "UCX_RNDV_THRESH" controls threshold threshold for switching between
+        # eager and rendezvous (RNDV) communication protocols.
+        # "UCX_NET_DEVICES" select which network interfaces UCX should use.
+        # "UCC_CL_BASIC_TLS" controls which Transport Layers are used by
+        # the Basic Collective libraray
+
+        os.environ['TORCH_UCC_BLOCKING_WAIT'] = (
+            os.environ['TORCH_UCC_BLOCKING_WAIT']
+            if "TORCH_UCC_BLOCKING_WAIT" in os.environ
+            else 'none'
+        )
+        os.environ['UCC_EC_CUDA_STREAM_TASK_MODE'] = (
+            os.environ['UCC_EC_CUDA_STREAM_TASK_MODE']
+            if "UCC_EC_CUDA_STREAM_TASK_MODE" in os.environ
+            else 'driver'
+        )
+        os.environ['UCX_TLS'] = (
+            os.environ['UCX_TLS'] if "UCX_TLS" in os.environ else 'ib,cuda_copy'
+        )  # cuda_ipc (i.e., NVLink-enablement) will be later supported
+        os.environ['NSYS_UCP_COMM_PARAMS'] = '1'
+        os.environ['UCX_RNDV_THRESH'] = '0'
+        os.environ['UCX_NET_DEVICES'] = 'all'
+        os.environ['UCC_CL_BASIC_TLS'] = '^sharp,nccl'
+
+    for ranks in generator_wrapper('pp'):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            backend=pipeline_model_parallel_comm_backend,
+            pg_options=(
+                None
+                if pipeline_model_parallel_comm_backend == 'ucc'
+                else get_nccl_options('pp', nccl_comm_cfgs)
+            ),
+            group_desc='PIPELINE_MODEL_PARALLEL_GROUP',
+        )
+        assert (
+            pipeline_model_parallel_comm_backend == None
+            or pipeline_model_parallel_comm_backend == 'nccl'
+            or pipeline_model_parallel_comm_backend == 'ucc'
+        ), f'"{pipeline_model_parallel_comm_backend}" backend for PP communication is currently not supported'
+
+        if rank in ranks:
+            if _PIPELINE_MODEL_PARALLEL_GROUP is None:
+                _PIPELINE_MODEL_PARALLEL_GROUP = group
+                _PIPELINE_GLOBAL_RANKS = ranks
+            elif isinstance(_PIPELINE_GLOBAL_RANKS[0], list):
+                _PIPELINE_MODEL_PARALLEL_GROUP.append(group)
+                _PIPELINE_GLOBAL_RANKS.append(ranks)
+            else:
+                _PIPELINE_MODEL_PARALLEL_GROUP = [_PIPELINE_MODEL_PARALLEL_GROUP, group]
+                _PIPELINE_GLOBAL_RANKS = [_PIPELINE_GLOBAL_RANKS, ranks]
+
+        embedding_ranks = get_embedding_ranks(ranks)
+        group = create_group(
+            embedding_ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('embd', nccl_comm_cfgs),
+            group_desc='EMBEDDING_GROUP',
+        )
+        if rank in embedding_ranks:
+            _EMBEDDING_GROUP = group
+            _EMBEDDING_GLOBAL_RANKS = embedding_ranks
+
+        position_embedding_ranks = get_position_embedding_ranks(ranks)
+        group = create_group(
+            position_embedding_ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('pos_embd', nccl_comm_cfgs),
+            group_desc='POSITION_EMBEDDING_GROUP',
+        )
+        if rank in position_embedding_ranks:
+            _POSITION_EMBEDDING_GROUP = group
+            _POSITION_EMBEDDING_GLOBAL_RANKS = position_embedding_ranks
+
+    # Build the tensor + data parallel groups.
+    global _TENSOR_AND_DATA_PARALLEL_GROUP
+    global _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP
+    assert (
+        _TENSOR_AND_DATA_PARALLEL_GROUP is None
+    ), 'Tensor + data parallel group is already initialized'
+    for ranks in generator_wrapper('tp-dp-cp'):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('tp_dp_cp', nccl_comm_cfgs),
+            group_desc='TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP',
+        )
+        if rank in ranks:
+            _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP = group
+    for ranks in generator_wrapper('tp-dp'):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('tp_dp', nccl_comm_cfgs),
+            group_desc='TENSOR_AND_DATA_PARALLEL_GROUP',
+        )
+        if rank in ranks:
+            _TENSOR_AND_DATA_PARALLEL_GROUP = group
+
+    global _TENSOR_AND_CONTEXT_PARALLEL_GROUP
+    assert (
+        _TENSOR_AND_CONTEXT_PARALLEL_GROUP is None
+    ), 'Tensor + context parallel group is already initialized'
+    for ranks in generator_wrapper('tp-cp'):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('tp_cp', nccl_comm_cfgs),
+            group_desc='TENSOR_AND_CONTEXT_PARALLEL_GROUP',
+        )
+        if rank in ranks:
+            _TENSOR_AND_CONTEXT_PARALLEL_GROUP = group
+
+    ### Expert-related parallel groups initialization
+    # Build the expert model parallel group
+    global _EXPERT_MODEL_PARALLEL_GROUP
+    assert _EXPERT_MODEL_PARALLEL_GROUP is None, 'Expert parallel group is already initialized'
+    for ranks in generator_wrapper('ep', is_expert=True):
+        group = create_group(
+            ranks,
+            pg_options=get_nccl_options('ep', nccl_comm_cfgs),
+            group_desc='EXPERT_MODEL_PARALLEL_GROUP',
+        )
+        if rank in ranks:
+            _EXPERT_MODEL_PARALLEL_GROUP = group
+
+    # Build the expert tensor parallel group
+    global _EXPERT_TENSOR_PARALLEL_GROUP
+    assert (
+        _EXPERT_TENSOR_PARALLEL_GROUP is None
+    ), 'Expert tensor model parallel group is already initialized'
+    for ranks in generator_wrapper('tp', is_expert=True):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('ep_tp', nccl_comm_cfgs),
+            group_desc='EXPERT_TENSOR_PARALLEL_GROUP',
+        )
+        if rank in ranks:
+            _EXPERT_TENSOR_PARALLEL_GROUP = group
+
+    # Build the tensor + expert parallel groups
+    global _EXPERT_TENSOR_AND_MODEL_PARALLEL_GROUP
+    assert (
+        _EXPERT_TENSOR_AND_MODEL_PARALLEL_GROUP is None
+    ), 'Expert tensor + model parallel group is already initialized'
+    for ranks in generator_wrapper('tp-ep', is_expert=True):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('tp_ep_mp', nccl_comm_cfgs),
+            group_desc='EXPERT_TENSOR_AND_MODEL_PARALLEL_GROUP',
+        )
+        if rank in ranks:
+            _EXPERT_TENSOR_AND_MODEL_PARALLEL_GROUP = group
+
+    # Build the expert+tensor+pipeline parallel groups
+    global _EXPERT_TENSOR_MODEL_PIPELINE_PARALLEL_GROUP
+    assert (
+        _EXPERT_TENSOR_MODEL_PIPELINE_PARALLEL_GROUP is None
+    ), 'The expert_tensor_model_pipeline parallel group is already initialized'
+    for ranks in generator_wrapper('tp-ep-pp', is_expert=True):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('tp_ep_pp', nccl_comm_cfgs),
+            group_desc='EXPERT_TENSOR_MODEL_PIPELINE_PARALLEL_GROUP',
+        )
+        if rank in ranks:
+            _EXPERT_TENSOR_MODEL_PIPELINE_PARALLEL_GROUP = group
+
+    # Build the expert data parallel group
+    global _EXPERT_DATA_PARALLEL_GROUP
+    assert _EXPERT_DATA_PARALLEL_GROUP is None, 'Expert data group is already initialized'
+    global _EXPERT_DATA_PARALLEL_GROUP_GLOO
+    assert _EXPERT_DATA_PARALLEL_GROUP_GLOO is None, 'Expert data group-gloo is already initialized'
+
+    for ranks in generator_wrapper('dp', is_expert=True):
+        group = create_group(
+            ranks,
+            timeout=timeout,
+            pg_options=get_nccl_options('ep_dp', nccl_comm_cfgs),
+            group_desc='EXPERT_DATA_PARALLEL_GROUP',
+        )
+        if create_gloo_process_groups:
+            group_gloo = create_group(
+                ranks, backend="gloo", group_desc='EXPERT_DATA_PARALLEL_GROUP_GLOO'
+            )
+        else:
+            group_gloo = None
+        if rank in ranks:
+            _EXPERT_DATA_PARALLEL_GROUP = group
+            _EXPERT_DATA_PARALLEL_GROUP_GLOO = group_gloo
+    ### End of expert related parallel groups initialization
+
+    # Initialize global memory buffer
+    # This isn't really "parallel state" but there isn't another good place to
+    # put this. If we end up with a more generic initialization of megatron-core
+    # we could stick it there
+    _set_global_memory_buffer()
+
+    for var in list(group_list.keys())[8:]:
+        setattr(sys.modules["megatron.core.parallel_state"], var, eval(var))
+
+def get_epx_data_parallel_lcp():
+        return parallel_state._EPX_DATA_PARALLEL_LCP
+
+# use for fault_tolerance
+# initialize_model_parallel only update to set _EPX_DATA_PARALLEL_LCP, and no other changes
+# get_epx_data_parallel_lcp used to get _EPX_DATA_PARALLEL_LCP.
+# _EPX_DATA_PARALLEL_LCP is only used in fault_tolerance
+attrs_to_register = ['initialize_model_parallel', 'get_epx_data_parallel_lcp']
+
+for k in sys.modules:
+    if k.endswith('megatron.core.parallel_state'):
+        for target in attrs_to_register:
+            setattr(sys.modules[k], target, eval(target))
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/profiling.py b/megatron-lm-musa-patch/musa_patch/profiling.py
new file mode 100644
index 000000000..0a38dbf4d
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/profiling.py
@@ -0,0 +1,160 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import contextlib
+import os
+import pickle
+import time
+from datetime import datetime
+from pathlib import Path
+import json
+
+import torch
+
+# the number of warmup steps before the active step in each profiling cycle
+profile_freq = 4
+# how much memory allocation/free ops to record in memory snapshots
+MEMORY_SNAPSHOT_MAX_ENTRIES = 100000
+
+
+@contextlib.contextmanager
+def maybe_enable_profiling(args, global_step):
+    #add tarce related centext: renll
+    on_demand_profiling = int(os.getenv("KINETO_USE_DAEMON", 0))
+    if on_demand_profiling == 1:
+        training_job_path = os.getenv("TRAINING_JOB_PATH", "/home/dist")
+        rank_pid_relation_dir = os.getenv("RANK_PID_RELATION_DIR", "rank_pid_relation_dir")
+        job_id = os.getenv("MCCFLOW_JOB_ID", "default-job-id")
+        rank_full_path = "{}/{}/{}".format(training_job_path, rank_pid_relation_dir, job_id)
+        Path(rank_full_path).mkdir(parents=True, exist_ok=True)
+
+        rank_pid_relation_map = {
+            "rank": torch.distributed.get_rank(),
+            "world_size": torch.distributed.get_world_size(),
+            "pid": os.getpid(),
+        }
+        rank_file = rank_full_path + "/rank" + str(torch.distributed.get_rank())
+        with open(rank_file, "w+") as f:
+            json.dump(rank_pid_relation_map, f)
+    # get user defined profiler settings
+    enable_profiling = int(os.getenv("ENABLE_PROFILER", 0))
+     # fetch profiler related env
+    wait_steps = int(os.getenv("PROFILER_WAIT_STEPS", 0))
+    warmup_steps = int(os.getenv("PROFILER_WARMUP_STEPS", 3))
+    active_steps = int(os.getenv("PROFILER_ACTIVE_STEPS", 1))
+    repeat_num = int(os.getenv("PROFILER_REPEAT_NUM", 0))
+    profile_freq = int(os.getenv("PROFILER_FREQ", 1))
+    current_time = datetime.now().strftime("%Y.%m.%d-%H:%M:%S")
+    save_dir = os.getenv("PROFILER_SAVE_DIR", f"./profiler_result/{current_time}")
+    worker_name = os.getenv(
+        "PROFILER_WORKER_NAME", "rank" + str(torch.distributed.get_rank())
+    )
+    record_shapes = int(os.getenv("PROFILER_RECORD_SHAPES", 1))
+    profile_memory = int(os.getenv("PROFILER_PROFILE_MEMORY", 0))
+    with_stack = int(os.getenv("PROFILER_WITH_STACK", 1))
+    with_modules = int(os.getenv("PROFILER_WITH_MODULES", 1))
+    kineto_log_level = int(os.getenv("KINETO_LOG_LEVEL", 0))
+
+    if enable_profiling:
+        profile_freq = profile_freq
+
+        rank = torch.distributed.get_rank()
+
+        def trace_handler(prof):
+            curr_trace_dir_name = "iteration_" + str(prof.step_num)
+            curr_trace_dir = os.path.join(save_dir, curr_trace_dir_name)
+            if not os.path.exists(curr_trace_dir):
+                os.makedirs(curr_trace_dir, exist_ok=True)
+            curr_trace_path = os.path.join(curr_trace_dir, f"rank{rank}.{int(time.time()*1000)}.pt.trace.json")
+            print(f"Dumping profiler traces at step {prof.step_num} to {curr_trace_path}")
+            begin = time.monotonic()
+            prof.export_chrome_trace(curr_trace_path)
+            print(
+                f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"
+            )
+
+        print(f"Profiling active. Traces will be saved at {save_dir}")
+
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir, exist_ok=True)
+
+        # warmup, active = WARMUP, 1
+        wait = profile_freq - (active_steps + warmup_steps)
+        assert (
+            wait >= 0
+        ), "profile_freq must be greater than or equal to warmup + active"
+        with torch.profiler.profile(
+            activities=[
+                torch.profiler.ProfilerActivity.CPU,
+                torch.profiler.ProfilerActivity.MUSA,
+                # torch.profiler.ProfilerActivity.CUDA,
+            ],
+            schedule=torch.profiler.schedule(wait=wait, warmup=warmup_steps, active=active_steps, repeat=repeat_num),
+            on_trace_ready=trace_handler,
+            record_shapes=record_shapes,
+            profile_memory=profile_memory,
+            with_stack=with_stack,
+            with_modules=with_modules,
+            start_step=global_step,
+        ) as torch_profiler:
+            yield torch_profiler
+    else:
+        torch_profiler = contextlib.nullcontext()
+        yield None
+
+
+@contextlib.contextmanager
+def maybe_enable_memory_snapshot(args, global_step: int = 0):
+    pass
+    # enable_snapshot = config.profiling.enable_memory_snapshot
+    # if enable_snapshot:
+    #     snapshot_folder = config.profiling.save_memory_snapshot_folder
+    #     snapshot_dir = os.path.join(config.job.dump_folder, snapshot_folder)
+    #     if not os.path.exists(snapshot_dir):
+    #         os.makedirs(snapshot_dir, exist_ok=True)
+    #     rank = torch.distributed.get_rank()
+
+    #     class MemoryProfiler:
+    #         def __init__(self, step_num: int, freq: int):
+    #             torch.musa.memory._record_memory_history(
+    #                 max_entries=MEMORY_SNAPSHOT_MAX_ENTRIES
+    #             )
+    #             # when resume training, we start from the last step
+    #             self.step_num = step_num
+    #             self.freq = freq
+
+    #         def step(self, exit_ctx: bool = False):
+    #             self.step_num += 1
+    #             if not exit_ctx and self.step_num % self.freq != 0:
+    #                 return
+    #             if not exit_ctx:
+    #                 curr_step = self.step_num
+    #                 dir_name = f"iteration_{curr_step}"
+    #             else:
+    #                 # dump as iteration_0_exit if OOM at iter 1
+    #                 curr_step = self.step_num - 1
+    #                 dir_name = f"iteration_{curr_step}_exit"
+    #             curr_snapshot_dir = os.path.join(snapshot_dir, dir_name)
+    #             if not os.path.exists(curr_snapshot_dir):
+    #                 os.makedirs(curr_snapshot_dir, exist_ok=True)
+    #             logger.info(f"Dumping memory snapshot at step {curr_step}")
+    #             begin = time.monotonic()
+    #             with open(
+    #                 f"{curr_snapshot_dir}/rank{rank}_memory_snapshot.pickle", "wb"
+    #             ) as output:
+    #                 pickle.dump(torch.musa.memory._snapshot(), output)
+    #             logger.info(
+    #                 f"Finished dumping memory snapshot in {time.monotonic() - begin:.2f} seconds"
+    #             )
+
+    #     logger.info(f"Memory profiler active. Snapshot will be saved at {snapshot_dir}")
+    #     profiler = MemoryProfiler(global_step, config.profiling.profile_freq)
+    #     try:
+    #         yield profiler
+    #     except torch.OutOfMemoryError as e:
+    #         profiler.step(exit_ctx=True)
+    # else:
+    #     yield None
diff --git a/megatron-lm-musa-patch/musa_patch/recomupte_variance/__init__.py b/megatron-lm-musa-patch/musa_patch/recomupte_variance/__init__.py
new file mode 100644
index 000000000..0437d1538
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/recomupte_variance/__init__.py
@@ -0,0 +1,7 @@
+from . import multi_latent_attention
+from . import random
+from . import moe_layer
+from . import mlp
+from . import transformer_layer
+from . import transformer_engine
+from . import schedules
diff --git a/megatron-lm-musa-patch/musa_patch/recomupte_variance/mlp.py b/megatron-lm-musa-patch/musa_patch/recomupte_variance/mlp.py
new file mode 100644
index 000000000..aa5f0c799
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/recomupte_variance/mlp.py
@@ -0,0 +1,144 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import torch
+import torch.nn.functional as F
+
+from megatron.core import tensor_parallel, parallel_state
+
+from megatron.core.fusions.fused_bias_geglu import bias_geglu_impl
+from megatron.core.fusions.fused_bias_gelu import bias_gelu_impl
+from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl
+
+from transformer_engine.pytorch.distributed import checkpoint, checkpointVirance
+# HACK(huang.huang): recompute/variance for mlp in moe with fp8/bf16:
+# support mlp_rms_recompute,  which combine rms, mlp into one checkpoint;
+# add new arg "no_recompute" to avoid repated recompute for sharedEXP while
+# moe_layer is already recomputed outsides
+def MLP_forward(self, hidden_states, norm_func=None, no_recompute=False):
+    """
+    Perform the forward pass through the MLP block.
+    Args:
+    hidden_states (Tensor): Input tensor of shape [s, b, h] where s is sequence length,
+        b is batch size, and h is hidden size.
+    norm_func (function): whether to do layernorm inner MLP instead of transformerlayer.
+    no_recompute (bool): default is False. only set to True when is sharedEXP,
+                        to avoid repeated recomputation between this mlp and moe_layer
+    """
+    # [s, b, 4 * h/p]
+    def custom_forward(hidden_states):
+        if norm_func is not None:
+            assert self.config.mlp_rms_recompute
+
+            hidden_states= norm_func(hidden_states)
+        intermediate_parallel, bias_parallel = self.linear_fc1(hidden_states)
+
+        if self.config.bias_activation_fusion:
+            if self.activation_func == F.gelu:
+                if self.config.gated_linear_unit:
+                    intermediate_parallel = bias_geglu_impl(intermediate_parallel, bias_parallel)
+                else:
+                    assert self.config.add_bias_linear is True
+                    intermediate_parallel = bias_gelu_impl(intermediate_parallel, bias_parallel)
+            elif self.activation_func == F.silu and self.config.gated_linear_unit:
+                intermediate_parallel = bias_swiglu_impl(
+                    intermediate_parallel,
+                    bias_parallel,
+                    self.config.activation_func_fp8_input_store,
+                )
+            else:
+                raise ValueError("Only support fusion of gelu and swiglu")
+        else:
+            if bias_parallel is not None:
+                intermediate_parallel = intermediate_parallel + bias_parallel
+            if self.config.gated_linear_unit:
+
+                def glu(x):
+                    x = torch.chunk(x, 2, dim=-1)
+                    return self.config.activation_func(x[0]) * x[1]
+
+                intermediate_parallel = glu(intermediate_parallel)
+            else:
+                intermediate_parallel = self.activation_func(intermediate_parallel)
+
+        # [s, b, h]
+        output, output_bias = self.linear_fc2(intermediate_parallel)
+        return output, output_bias
+
+    if norm_func is not None:
+        _custom_func_first = lambda x : self.custom_func_first(norm_func(x))
+    else:
+        _custom_func_first = lambda x : self.custom_func_first(x)# use lambda to create new func instead of method object which can't add new attribute
+    if no_recompute: #avoid to recompute under another recompute context outside this function, like in sharedExp
+        return custom_forward(hidden_states)
+
+    if self.config.mlp_recompute:
+        if self.config.fp8:
+            if self.config.recompute_variance:
+                output, output_bias = checkpointVirance(
+                    _custom_func_first,
+                    self.linear_fc2,
+                    hidden_states,
+                    distribute_saved_activations=self.config.distribute_saved_activations,
+                    get_rng_state_tracker=tensor_parallel.random.get_cuda_rng_tracker,
+                    tp_group=parallel_state.get_tensor_model_parallel_group(),
+                )
+            else:
+                output, output_bias = checkpoint(
+                    custom_forward,
+                    hidden_states,
+                    distribute_saved_activations=self.config.distribute_saved_activations,
+                    get_rng_state_tracker=tensor_parallel.random.get_cuda_rng_tracker,
+                    tp_group=parallel_state.get_tensor_model_parallel_group(),
+                    )
+        else:
+            if self.config.recompute_variance:
+                output, output_bias = tensor_parallel.checkpointVirance(
+                    _custom_func_first, self.linear_fc2, False, hidden_states)
+            else:
+                output, output_bias = tensor_parallel.checkpoint(
+                    custom_forward, False, hidden_states)
+    else:
+        output, output_bias = custom_forward(hidden_states)
+    return output, output_bias
+## HACK(huang.huang)
+
+# HACK(huang.huang): seperate linear1 and act from mlp, to support potential recoumpute variance,
+# which need a separated linear2
+def MLP_custom_func_first(self, hidden_states):
+    intermediate_parallel, bias_parallel = self.linear_fc1(hidden_states)
+
+    if self.config.bias_activation_fusion:
+        if self.activation_func == F.gelu:
+            if self.config.gated_linear_unit:
+                intermediate_parallel = bias_geglu_impl(intermediate_parallel, bias_parallel)
+            else:
+                assert self.config.add_bias_linear is True
+                intermediate_parallel = bias_gelu_impl(intermediate_parallel, bias_parallel)
+        elif self.activation_func == F.silu and self.config.gated_linear_unit:
+            intermediate_parallel = bias_swiglu_impl(
+                intermediate_parallel,
+                bias_parallel,
+                self.config.activation_func_fp8_input_store,
+            )
+        else:
+            raise ValueError("Only support fusion of gelu and swiglu")
+    else:
+        if bias_parallel is not None:
+            intermediate_parallel = intermediate_parallel + bias_parallel
+        if self.config.gated_linear_unit:
+
+            def glu(x):
+                x = torch.chunk(x, 2, dim=-1)
+                return self.config.activation_func(x[0]) * x[1]
+
+            intermediate_parallel = glu(intermediate_parallel)
+        else:
+            intermediate_parallel = self.activation_func(intermediate_parallel)
+
+    return intermediate_parallel
+## HACK(huang.huang)
+
+from transformer_engine.musa.pytorch.utils import replace_attr, add_attr
+from megatron.core.transformer.mlp import MLP
+replace_attr(MLP,"forward", MLP_forward)
+add_attr(MLP,"custom_func_first", MLP_custom_func_first)
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/recomupte_variance/moe_layer.py b/megatron-lm-musa-patch/musa_patch/recomupte_variance/moe_layer.py
new file mode 100644
index 000000000..ad0b0b979
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/recomupte_variance/moe_layer.py
@@ -0,0 +1,216 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from functools import partial, wraps
+
+import torch
+import torch.nn.functional as F
+from torch.nn import Identity
+
+from megatron.core import tensor_parallel, parallel_state
+from megatron.core.fusions.fused_bias_geglu import bias_geglu_impl
+from megatron.core.fusions.fused_bias_gelu import bias_gelu_impl
+from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl
+
+from transformer_engine.pytorch.distributed import checkpoint, checkpointVirance
+from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
+
+# HACK(huang.huang): recompute/variance for experts in moe with fp8/bf16:
+# support mlp_rms_recompute which combine rms, sharedEXP and gating into one checkpoint;
+def MoELayer_forward(self, hidden_states: torch.Tensor, norm_func=None):
+    """
+    Perform the forward pass through the MLP block.
+    Args:
+    hidden_states (Tensor): Input tensor of shape [s, b, h] where s is sequence length,
+        b is batch size, and h is hidden size.
+    norm_func (function): whether to do layernorm inner MLP instead of transformerlayer.
+    """
+    if (
+        self.training
+        and self.config.tensor_model_parallel_size > 1
+        and not self.config.sequence_parallel
+    ):
+        raise ValueError(
+            "During training, performance may degrade if MoE and tensor parallelism"
+            "are enabled without also enabling sequence parallelism."
+        )
+
+
+    # process MoE
+    def custom_forward(hidden_states):
+
+        if norm_func is not None:
+            assert self.config.mlp_rms_recompute
+
+            def rms_recompute_func(hidden_states):
+                #combination of rms, sharedEXP and gating
+                hidden_states= norm_func(hidden_states)
+                logits = self.router.apply_input_jitter(hidden_states)
+                logits = self.router.gating(logits)
+                shared_output = self.shared_experts(hidden_states, no_recompute=True)
+                return hidden_states, logits, shared_output
+
+            if self.config.fp8:
+                if self.config.recompute_variance:
+                    func_before_routing = lambda x : self.router.gating(self.router.apply_input_jitter(x))
+                    linears = (Identity(), Identity(), self.shared_experts.linear_fc2)
+                    mid_function = (Identity(), func_before_routing, self.shared_experts.custom_func_first)
+                    hidden_states, logits, shared_output = checkpointVirance(
+                        norm_func,
+                        linears,
+                        hidden_states,
+                        distribute_saved_activations=self.config.distribute_saved_activations,
+                        get_rng_state_tracker=tensor_parallel.random.get_cuda_rng_tracker,
+                        tp_group=parallel_state.get_tensor_model_parallel_group(),
+                        mid_function=mid_function
+                    )
+                else:
+                    hidden_states, logits, shared_output = checkpoint(
+                        rms_recompute_func,
+                        hidden_states,
+                        distribute_saved_activations=self.config.distribute_saved_activations,
+                        get_rng_state_tracker=tensor_parallel.random.get_cuda_rng_tracker,
+                        tp_group=parallel_state.get_tensor_model_parallel_group(),
+                    )
+            else:
+                if self.config.recompute_variance:
+                    func_before_routing = lambda x : self.router.gating(self.router.apply_input_jitter(x))
+                    linears = (Identity(), Identity(), self.shared_experts.linear_fc2)
+                    mid_function = (Identity(), func_before_routing, self.shared_experts.custom_func_first)
+                    hidden_states, logits, shared_output = tensor_parallel.checkpointVirance(
+                        norm_func,
+                        linears,
+                        False,
+                        hidden_states,
+                        mid_function=mid_function
+                        )
+                else:
+                    hidden_states, logits, shared_output = tensor_parallel.checkpoint(
+                        rms_recompute_func, False, hidden_states)
+            probs, routing_map = self.router.routing(logits)
+        else:
+            probs, routing_map = self.router(hidden_states)
+        (dispatched_input, tokens_per_expert) = self.token_dispatcher.token_permutation(
+            hidden_states, probs, routing_map
+        )
+        custom_expert_forward = partial(self.experts, tokens_per_expert=tokens_per_expert)
+
+        def _custom_func_first(permuted_local_hidden_states, tokens_per_expert):
+            #forward for linear1 and act in self.experts
+            tokens_per_expert = tokens_per_expert.tolist()
+            intermediate_parallel, bias_parallel = self.experts.linear_fc1(
+                permuted_local_hidden_states, tokens_per_expert
+            )
+
+            if self.experts.config.bias_activation_fusion:
+                if self.experts.activation_func == F.gelu:
+                    if self.experts.config.gated_linear_unit:
+                        intermediate_parallel = bias_geglu_impl(intermediate_parallel, bias_parallel)
+                    else:
+                        assert self.experts.config.add_bias_linear is True
+                        intermediate_parallel = bias_gelu_impl(intermediate_parallel, bias_parallel)
+                elif self.experts.activation_func == F.silu and self.experts.config.gated_linear_unit:
+                    intermediate_parallel = bias_swiglu_impl(
+                        intermediate_parallel,
+                        bias_parallel,
+                        self.config.activation_func_fp8_input_store,
+                    )
+                else:
+                    raise ValueError("Only support fusion of gelu and swiglu")
+            else:
+                if bias_parallel is not None:
+                    shape = intermediate_parallel.shape
+                    intermediate_parallel = torch.cat(
+                        [
+                            t + b
+                            for t, b in zip(
+                                torch.split(
+                                    intermediate_parallel.view(-1, shape[-1]), tokens_per_expert
+                                ),
+                                bias_parallel,
+                            )
+                        ]
+                    ).view(shape)
+                if self.experts.config.gated_linear_unit:
+
+                    def glu(x):
+                        x = torch.chunk(x, 2, dim=-1)
+                        return self.experts.config.activation_func(x[0]) * x[1]
+
+                    intermediate_parallel = glu(intermediate_parallel)
+                else:
+                    intermediate_parallel = self.experts.activation_func(intermediate_parallel)
+            return intermediate_parallel, tokens_per_expert
+
+        custom_func_first = partial(_custom_func_first, tokens_per_expert=tokens_per_expert)
+
+
+        if self.config.mlp_recompute:
+            if self.config.fp8:
+                if self.config.recompute_variance:
+                    expert_output, mlp_bias = checkpointVirance(
+                        custom_func_first,
+                        self.experts.linear_fc2,
+                        dispatched_input,
+                        distribute_saved_activations=self.config.distribute_saved_activations,
+                        get_rng_state_tracker=tensor_parallel.random.get_cuda_rng_tracker,
+                        tp_group=parallel_state.get_tensor_model_parallel_group(),
+                    )
+                else:
+                    expert_output, mlp_bias = checkpoint(
+                        custom_expert_forward,
+                        dispatched_input,
+                        distribute_saved_activations=self.config.distribute_saved_activations,
+                        get_rng_state_tracker=tensor_parallel.random.get_cuda_rng_tracker,
+                        tp_group=parallel_state.get_tensor_model_parallel_group(),
+                        )
+            else:
+                if self.config.recompute_variance:
+                    expert_output, mlp_bias = tensor_parallel.checkpointVirance(
+                        custom_func_first, self.experts.linear_fc2, False, dispatched_input)
+                else:
+                    expert_output, mlp_bias = tensor_parallel.checkpoint(
+                        custom_expert_forward, False, dispatched_input)
+        else:
+            expert_output, mlp_bias = self.experts(dispatched_input, tokens_per_expert)
+
+        output, mlp_bias = self.token_dispatcher.token_unpermutation(expert_output, mlp_bias)
+        if norm_func is not None:
+            #self.shared_experts called in the begining of custom_forward, which is convenient for rms recmopute
+            output = output + shared_output
+        elif self.use_shared_expert and not self.shared_expert_overlap:
+            # if shared_expert_overlap is True, the expert calculation happens in
+            # the token_dispatcher to overlap communications and computations
+            output = output + self.shared_experts(hidden_states)
+        return output, mlp_bias
+
+    if self.moe_layer_recompute:
+        output, mlp_bias = tensor_parallel.checkpoint(custom_forward, False, hidden_states)
+    else:
+        output, mlp_bias = custom_forward(hidden_states)
+
+    return output, mlp_bias
+## HACK(huang.huang)
+
+
+# HACK(huang.huang): recompute/variance for SharedExpertMLP, avoid repeated recomputation between moe-layer and sharedExp
+def SharedExpertMLP_forward(self, hidden_states, no_recompute=False):
+    """
+    Perform the forward pass through the SharedExpertMLP block.
+    Args:
+    hidden_states (Tensor): Input tensor of shape [s, b, h] where s is sequence length,
+        b is batch size, and h is hidden size.
+    no_recompute (bool): default is False. only set to True when is sharedEXP,
+                        to avoid repeated recomputation between this mlp and moe_layer
+    """
+    output, _ = super(SharedExpertMLP, self).forward(hidden_states, no_recompute=no_recompute)
+    if self.use_shared_expert_gate:
+        logits = torch.nn.functional.linear(hidden_states, self.gate_weight)
+        gate_score = torch.nn.functional.sigmoid(logits)
+        output = output * gate_score
+    return output
+## HACK(huang.huang)
+
+from transformer_engine.musa.pytorch.utils import replace_attr
+from megatron.core.transformer.moe.moe_layer import MoELayer
+replace_attr(MoELayer,"forward", MoELayer_forward)
+replace_attr(SharedExpertMLP,"forward", SharedExpertMLP_forward)
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/recomupte_variance/multi_latent_attention.py b/megatron-lm-musa-patch/musa_patch/recomupte_variance/multi_latent_attention.py
new file mode 100644
index 000000000..f1068c09a
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/recomupte_variance/multi_latent_attention.py
@@ -0,0 +1,297 @@
+import os
+import torch
+from functools import partial
+from megatron.core import tensor_parallel, parallel_state
+from megatron.core.models.common.embeddings import apply_rotary_pos_emb
+from megatron.core.transformer.multi_latent_attention import MLASelfAttention
+try:
+    from transformer_engine.pytorch.distributed import checkpoint
+    from transformer_engine.pytorch.distributed import checkpointViranceAttention
+    HAVE_TE = True
+except ImportError:
+    HAVE_TE = False
+
+# HACK(huang.huang): recompute-variance for fa:
+# 1. modify get_query_key_value_tensors for MLASelfAttention, just add a logic to call recompute;
+# 2. modify forward for MLASelfAttention, seperate the core attention from other part around it, and send them to checkpoint_forward
+# 3. add RoPEQInplace
+# TODO: huang.huang revise code before to follow new version "get_qkv" in Megatron-LM:main
+class RoPEQInplace(torch.autograd.Function):
+    """
+    limiation:
+    1. pre_op backward cannot use self output(e.g. softmax).
+    2. if you call backward directly and pass in dy, be careful that dy is overwritten.
+    """
+
+    @staticmethod
+    def forward(ctx, x, freqs, custom_metadata):
+        (
+            split_start,
+            split_end,
+            rotary_interleaved,
+            batch_first,
+        ) = ctx.custom_metadata = custom_metadata
+        assert x.dim() == 4 and freqs.dim() == 2
+        assert (split_end - split_start) == freqs.shape[-1]
+        assert x.shape[batch_first] == freqs.shape[0]
+        ctx.save_for_backward(freqs)
+        y = torch.ops.aten._fused_rope_forward(
+            x[..., split_start:split_end], freqs, rotary_interleaved, batch_first
+        )
+        # x.data[..., split_start:split_end] = y # Using `tensor.data` does not affect `tensor._version`.
+        x[..., split_start:split_end] = y
+        return x
+
+    @staticmethod
+    def backward(ctx, dy):
+        (freqs,) = ctx.saved_tensors
+        (
+            split_start,
+            split_end,
+            rotary_interleaved,
+            batch_first,
+        ) = ctx.custom_metadata
+        sub_dy = dy[..., split_start:split_end]
+        dx = torch.ops.aten._fused_rope_backward(
+            sub_dy, freqs, rotary_interleaved, batch_first
+        )
+        dy[..., split_start:split_end] = dx
+        return dy, None, None
+
+
+def MLASelfAttention_forward(
+    self,
+    hidden_states,
+    attention_mask,
+    key_value_states=None,
+    inference_params=None,
+    rotary_pos_emb=None,
+    rotary_pos_cos=None,
+    rotary_pos_sin=None,
+    attention_bias=None,
+    packed_seq_params=None,
+    position_ids=None,
+    sequence_len_offset=None,
+    q_compressed=None,
+    kv_combined=None,
+):
+    if not int(os.getenv("USE_RECOMPUTE_VARIANCE", 0)):
+        #original forward
+        return super(MLASelfAttention ,self).forward(
+            hidden_states,
+            attention_mask,
+            key_value_states,
+            inference_params,
+            rotary_pos_emb,
+            rotary_pos_cos,
+            rotary_pos_sin,
+            attention_bias,
+            packed_seq_params,
+            position_ids,
+            sequence_len_offset
+            )
+
+    """Forward pass for multi-latent attention"""
+    assert rotary_pos_emb is None, "Rotary position embeddings should not be passed into MLA."
+    assert attention_bias is None, "Attention bias should not be passed into MLA."
+    assert (
+        rotary_pos_cos is None and rotary_pos_sin is None
+    ), "MLA does not support Flash Decoding"
+
+    # hidden_states: [sq, b, h]
+
+    # =====================
+    # Query, Key, and Value
+    # =====================
+    # Get the query, key and value tensors based on the type of attention -
+    # self or cross attn.
+    # query: [96, 1, 16, 128], key:[96, 1, 16, 128], value:[96, 1, 16, 128]
+    if self.config.mla_rms_recompute:
+        assert self.config.attn_recompute, 'mla_rms_recompute only use with attn_recompute now.'
+        pass
+    else:
+        assert (
+            hidden_states.ndim == 3
+        ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D"
+
+        if self.config.q_lora_rank is not None:
+            q_compressed, _ = self.linear_q_down_proj(hidden_states)
+        else:
+            q_compressed = hidden_states
+
+        kv_combined, _ = self.linear_kv_down_proj(hidden_states)
+
+    def _custom_forward_before_attention(
+        q_compressed,
+        kv_combined,
+        key_value_states=None,
+        position_ids=None,
+        packed_seq_params=None,
+        inference_params=None,
+    ):
+        q_len, bsz, _ = q_compressed.size()
+
+        if self.config.q_lora_rank is not None:
+            q_compressed = self.q_layernorm(q_compressed)
+            q, _ = self.linear_q_up_proj(q_compressed)
+        else:
+            q, _ = self.linear_q_proj(q_compressed)
+
+        # q: [s, b, n, 192]
+        q = q.view(q_len, bsz, self.num_attention_heads_per_partition, self.q_head_dim)
+
+        # q: [s, b, n, 128], q_pos_emb: [s, b, n, 64]
+        q_no_pe, q_pos_emb = torch.split(
+            q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1
+        )
+
+
+        # kv_compressed:[s, b, 512], k_pos_emb: [s, b, 64]
+        kv_compressed, k_pos_emb = torch.split(
+            kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1
+        )
+        kv, _ = self.linear_kv_up_proj(self.kv_layernorm(kv_compressed))
+
+        # kv: [s, b, n, 256]
+        kv = kv.view(
+            q_len,
+            bsz,
+            self.num_attention_heads_per_partition,
+            self.config.qk_head_dim + self.config.v_head_dim,
+        )
+
+        # k_no_pe: [s, b, n, 128], value: [s, b, n, 128]
+        k_no_pe, value = torch.split(kv, [self.config.qk_head_dim, self.config.v_head_dim], dim=-1)
+
+        # rotary_pos_emb:[s, b, 1, 64]
+        rotary_pos_emb = self.rotary_pos_emb(max_seq_len=self.config.max_position_embeddings)
+
+        if len(rotary_pos_emb) == 2:
+            mscale = rotary_pos_emb[1]
+            rotary_pos_emb = rotary_pos_emb[0]
+
+        if inference_params is not None:
+            # add offset to the sequence start for inference
+            sequence_start = inference_params.sequence_len_offset
+            sequence_end = sequence_start + q_len
+            rotary_pos_emb = rotary_pos_emb[sequence_start:sequence_end]
+
+        # [s, b, 64] -> [s, b, 1, 64]
+        k_pos_emb = torch.unsqueeze(k_pos_emb, 2)
+
+        if packed_seq_params is not None:
+            cu_seqlens_q = packed_seq_params.cu_seqlens_q
+            cu_seqlens_kv = packed_seq_params.cu_seqlens_kv
+        else:
+            cu_seqlens_q = cu_seqlens_kv = None
+
+        # # q_pos_emb: [s, b, n, 64], k_pos_emb:[s, b, 1, 64]
+        # q_pos_emb = apply_rotary_pos_emb(
+        #     q_pos_emb, rotary_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q, mscale=mscale
+        # )
+        k_pos_emb = apply_rotary_pos_emb(
+            k_pos_emb, rotary_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv, mscale=mscale
+        )
+
+        # # query: [s, b, n, 192]
+        # query = torch.cat([q_no_pe, q_pos_emb], dim=-1)
+        q_split_start = self.config.qk_head_dim
+        q_split_end = q_split_start + self.config.qk_pos_emb_head_dim
+        rotary_interleaved = False
+        batch_first = False
+        query = RoPEQInplace.apply(q, rotary_pos_emb.squeeze(1).squeeze(1),
+                                   (q_split_start, q_split_end, rotary_interleaved, batch_first))
+
+        # key: [s, b, n, 192]
+        k_pos_emb = k_pos_emb.expand(-1, -1, self.config.num_attention_heads, -1)
+        key = torch.cat([k_no_pe, k_pos_emb], dim=-1)
+
+        query = query.contiguous()
+        key = key.contiguous()
+        value = value.contiguous()
+
+        query, key, value, _, attn_mask_type = self._adjust_key_value_for_inference(
+        inference_params, query, key, value, rotary_pos_emb=None
+        )
+        return query, key, value, attention_mask, \
+            {"attn_mask_type":attn_mask_type, "attention_bias":attention_bias, "packed_seq_params":packed_seq_params}
+
+    def _custom_forward_self_attention(
+        q_compressed,
+        kv_combined,
+        key_value_states=None,
+        position_ids=None,
+        packed_seq_params=None,
+        inference_params=None,
+        ):
+
+        query, key, value, attention_mask, kwargs = _custom_forward_before_attention(q_compressed, kv_combined, key_value_states, position_ids, packed_seq_params,inference_params)
+        core_attn_out = self.core_attention(query, key, value, attention_mask, **kwargs)
+        return core_attn_out
+
+    custom_forward_self_attention = partial(
+        _custom_forward_self_attention,
+        key_value_states=key_value_states,
+        inference_params=inference_params,
+        position_ids=position_ids,
+        packed_seq_params=packed_seq_params,
+    )
+
+
+    custom_forward_before_attention = partial(
+        _custom_forward_before_attention,
+        key_value_states=key_value_states,
+        inference_params=inference_params,
+        position_ids=position_ids,
+        packed_seq_params=packed_seq_params,
+    )
+
+    if self.config.attn_recompute == True:
+        if self.config.fp8:
+            if self.config.recompute_variance == True:
+                core_attn_out = checkpointViranceAttention(
+                    custom_forward_before_attention,
+                    self.core_attention,
+                    q_compressed,
+                    kv_combined,
+                    distribute_saved_activations=self.config.distribute_saved_activations,
+                    get_rng_state_tracker=tensor_parallel.random.get_cuda_rng_tracker,
+                    tp_group=parallel_state.get_tensor_model_parallel_group(),
+                )
+            else:
+                core_attn_out =  checkpoint(
+                    custom_forward_self_attention,
+                    q_compressed,
+                    kv_combined,
+                    distribute_saved_activations=self.config.distribute_saved_activations,
+                    get_rng_state_tracker=tensor_parallel.random.get_cuda_rng_tracker,
+                    tp_group=parallel_state.get_tensor_model_parallel_group(),
+                )
+        else:
+            if self.config.recompute_variance == True:
+                core_attn_out = tensor_parallel.checkpointViranceAttention(
+                    custom_forward_before_attention, self.core_attention, False, q_compressed, kv_combined)
+            else:
+                core_attn_out = tensor_parallel.checkpoint(
+                    custom_forward_self_attention, False, q_compressed, kv_combined)
+    else:
+        core_attn_out = custom_forward_self_attention(q_compressed, kv_combined)
+    if packed_seq_params is not None:
+        # reshape to same output shape as unpacked case
+        # (t, np, hn) -> (t, b=1, h=np*hn)
+        # t is the pack size = sum (sq_i)
+        # note that batch is a dummy dimension in the packed case
+        core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1)
+
+    # =================
+    # Output. [sq, b, h]
+    # =================
+    output, bias = self.linear_proj(core_attn_out)
+
+    return output, bias
+# HACK(huang.huang)
+
+
+from transformer_engine.musa.pytorch.utils import replace_attr
+
+replace_attr(MLASelfAttention, "forward", MLASelfAttention_forward)
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/recomupte_variance/random.py b/megatron-lm-musa-patch/musa_patch/recomupte_variance/random.py
new file mode 100644
index 000000000..5fc498b8a
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/recomupte_variance/random.py
@@ -0,0 +1,309 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+# Parts of the code here are adapted from PyTorch
+# repo: https://github.com/pytorch/pytorch
+
+# import contextlib
+# import logging
+
+import torch
+# from torch import _C
+from torch.cuda import _lazy_call
+from torch.nn import Identity
+# from torch.cuda import device as device_ctx_manager
+from torch.utils.checkpoint import detach_variable
+
+# from megatron.core.parallel_state import (
+#     get_expert_model_parallel_rank,
+#     get_expert_tensor_parallel_rank,
+#     get_tensor_model_parallel_rank,
+# )
+from megatron.core.utils import is_te_min_version, safely_set_viewless_tensor_data
+
+from megatron.core.tensor_parallel.utils import gather_split_1d_tensor, split_tensor_into_1d_equal_chunks
+
+
+from megatron.core.tensor_parallel.random import (CheckpointFunction, get_cuda_rng_tracker,
+                                                   _set_cuda_rng_state)
+
+
+# HACK(huang.huang): recompute-variance for [somefunc+fa] and [somefunc+linear],
+# which can save a forward for fa/linear when backward recompute
+# 2025.4.2: support list of linear as last_function, and args "mid_function" to support complex situations
+class IdentityTupleOp(torch.nn.Module):
+    """
+    This is a placeholder for IdentityTupleOp(*args) -> args,
+    """
+
+    def __init__(self,):
+        super().__init__()
+
+    def forward(self, *args):
+        return args
+
+
+class CheckpointFunctionVirance(CheckpointFunction):
+    """Checkpoint Function
+
+    This function is adapted from torch.utils.checkpoint with two main changes:
+    1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state`
+    2) the states in the model parallel tracker are also properly tracked/set/reset.
+    """
+
+    # pylint: disable=missing-function-docstring
+    @staticmethod
+    def forward(ctx, run_function, last_function, mid_function, distribute_saved_activations, *args):
+        """Forward pass."""
+        if not isinstance(last_function, tuple):
+            last_function = (last_function, )
+        mid_function = tuple(IdentityTupleOp() for _ in last_function) if mid_function is None else mid_function
+        ctx.run_function = run_function
+        ctx.last_function = last_function
+        ctx.mid_function = mid_function
+        ctx.distribute_saved_activations = distribute_saved_activations
+
+        # Copy the rng states.
+        ctx.fwd_cpu_rng_state = torch.get_rng_state()
+        ctx.fwd_cuda_rng_state = torch.cuda.get_rng_state()
+        ctx.fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()
+
+        with torch.no_grad():
+            outputs = run_function(*args)
+            outputs = outputs if isinstance(outputs, tuple) else (outputs, )
+            total_outputs = []
+            for i, func in enumerate(last_function):
+                outputs_f = mid_function[i](*outputs)
+                outputs_f = outputs_f if isinstance(outputs_f, tuple) else (outputs_f, )
+                outputs_f = func(*outputs_f)
+                total_outputs.append(outputs_f)
+            if len(total_outputs)==1:
+                #maintain original behavior when only one last_function
+                total_outputs=total_outputs[0]
+            else:
+                flat_outputs = []
+                for outputs_f in total_outputs:
+                    if isinstance(outputs_f, tuple):
+                        #Manually remove bias_out which is 'None', and assign 'None' to grad-bias in the corresponding backward direction
+                        outputs_f = tuple([x for x in outputs_f if x is not None])
+                    flat_outputs.append(outputs_f)
+                total_outputs = flat_outputs
+                #The reentrant version does not consider tensors in nested structures (e.g., custom objects, lists, dicts, etc)
+                # as participating in autograd, while the non-reentrant version does
+                total_outputs = sum( [x if isinstance(x, tuple) else (x,) for x in total_outputs ], tuple())
+        # Divide hidden states across model parallel group and only keep
+        # the chunk corresponding to the current rank.
+        if distribute_saved_activations:
+            ctx.input_0_shape = args[0].data.shape
+            safely_set_viewless_tensor_data(
+                args[0], split_tensor_into_1d_equal_chunks(args[0].data, new_buffer=True)
+            )
+
+        # Store everything.
+        ctx.inputs = [arg if not torch.is_tensor(arg) else None for arg in args]
+        tensor_inputs = [arg if torch.is_tensor(arg) else None for arg in args]
+        ctx.save_for_backward(*tensor_inputs)
+
+        return total_outputs
+
+    # pylint: disable=missing-function-docstring
+    @staticmethod
+    def backward(ctx, *args):
+        """Backward pass."""
+        if not torch.autograd._is_checkpoint_valid():
+            raise RuntimeError(
+                "Checkpointing is not compatible with .grad(), "
+                "please use .backward() if possible"
+            )
+        # inputs = ctx.saved_tensors
+        inputs = tuple(
+            t if t is not None else arg for (t, arg) in zip(ctx.saved_tensors, ctx.inputs)
+        )
+        if ctx.distribute_saved_activations:
+            safely_set_viewless_tensor_data(
+                inputs[0], gather_split_1d_tensor(inputs[0].data).view(ctx.input_0_shape)
+            )
+
+        # Store the current states.
+        bwd_cpu_rng_state = torch.get_rng_state()
+        bwd_cuda_rng_state = torch.cuda.get_rng_state()
+        bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()
+
+        # Set the states to what it used to be before the forward pass.
+        torch.set_rng_state(ctx.fwd_cpu_rng_state)
+        _set_cuda_rng_state(ctx.fwd_cuda_rng_state)
+        get_cuda_rng_tracker().set_states(ctx.fwd_cuda_rng_state_tracker)
+
+        # Compute the forward pass.
+        detached_inputs = detach_variable(inputs)
+        with torch.enable_grad():
+            outputs = ctx.run_function(*detached_inputs)
+            outputs = outputs if isinstance(outputs, tuple) else (outputs, )
+            total_outputs = []
+            for i,func in enumerate(ctx.mid_function):
+                outputs_f = func(*outputs)
+                if isinstance(outputs_f, torch.Tensor):
+                    outputs_f = [outputs_f,]
+                total_outputs.append(outputs_f)
+        # Set the states back to what it was at the start of this function.
+        torch.set_rng_state(bwd_cpu_rng_state)
+        _set_cuda_rng_state(bwd_cuda_rng_state)
+        get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker)
+
+
+        total_grad_input = []
+        for i,func in enumerate(ctx.last_function):
+            if isinstance(func, Identity):
+                grad_input_f = args[i]
+            else:
+                # Assign 'None' to grad_bias to correspond to the operation of removing 'none' during forward
+                grad_out_bias = args[i] if isinstance(args[i], tuple) else (args[i], None)
+                grad_input_f = func.backward_custom(*total_outputs[i], *grad_out_bias)
+            if isinstance(grad_input_f, torch.Tensor):
+                grad_input_f = (grad_input_f,)
+            total_grad_input.append(grad_input_f)
+
+        total_outputs_with_grad = []
+        total_args_with_grad = []
+        for j, outputs in enumerate(total_outputs):
+            outputs_with_grad = []
+            args_with_grad = []
+            for i, output in enumerate(outputs):
+                if torch.is_tensor(output) and output.requires_grad:
+                    outputs_with_grad.append(output)
+                    args_with_grad.append(total_grad_input[j][i])
+            total_outputs_with_grad += outputs_with_grad
+            total_args_with_grad += args_with_grad
+        torch.autograd.backward(total_outputs_with_grad, total_args_with_grad)
+        grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs)
+        return (None, None, None, None) + grads
+
+def checkpointVirance(run_function, last_function, distribute_saved_activations, *args, mid_function=None):
+    """Checkpoint a model or part of the model.
+    This has been directly copied from torch.utils.checkpoint."""
+    return CheckpointFunctionVirance.apply(run_function, last_function, mid_function, distribute_saved_activations, *args)
+
+
+
+class CheckpointFunctionViranceAttention(CheckpointFunction):
+    """Checkpoint Function
+
+    This function is adapted from torch.utils.checkpoint with two main changes:
+    1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state`
+    2) the states in the model parallel tracker are also properly tracked/set/reset.
+    """
+
+    # pylint: disable=missing-function-docstring
+    @staticmethod
+    def forward(ctx, run_function, last_function, distribute_saved_activations, *args):
+        """Forward pass."""
+        ctx.run_function = run_function
+        ctx.last_function = last_function
+        ctx.distribute_saved_activations = distribute_saved_activations
+
+        # Copy the rng states.
+        ctx.fwd_cpu_rng_state = torch.get_rng_state()
+        ctx.fwd_cuda_rng_state = torch.cuda.get_rng_state()
+        ctx.fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()
+
+        with torch.no_grad():
+            outputs = run_function(*args)
+            outputs = last_function.forward_before_fa(*outputs[:4], **outputs[4])
+            outputs = last_function.forward_fa(*outputs)
+            #outputs: Union[output=Union[Tensor output, Tensor logsumexp, Tensor dropout_mask],
+            # qkv_format, indices_q, batch_size, attn_mask_type, max_seqlen_q, q_shape, v_shape]
+            core_attn_out = last_function.forward_after_fa(*outputs)
+        # Divide hidden states across model parallel group and only keep
+        # the chunk corresponding to the current rank.
+        if distribute_saved_activations:
+            ctx.input_0_shape = args[0].data.shape
+            safely_set_viewless_tensor_data(
+                args[0], split_tensor_into_1d_equal_chunks(args[0].data, new_buffer=True)
+            )
+
+        # Store everything.
+        ctx.save_for_backward(*args, *outputs[0])
+        (ctx.qkv_format, ctx.indices_q, ctx.batch_size,
+         ctx.attn_mask_type, ctx.max_seqlen_q, ctx.q_shape, ctx.v_shape) = outputs[1:]
+
+        return core_attn_out
+
+# pylint: disable=missing-function-docstring
+    @staticmethod
+    def backward(ctx, *args):
+        """Backward pass."""
+        if not torch.autograd._is_checkpoint_valid():
+            raise RuntimeError(
+                "Checkpointing is not compatible with .grad(), "
+                "please use .backward() if possible"
+            )
+        inputs = ctx.saved_tensors
+        fa_output = inputs[-3:]
+        inputs = inputs[:-3]
+        if ctx.distribute_saved_activations:
+            safely_set_viewless_tensor_data(
+                inputs[0], gather_split_1d_tensor(inputs[0].data).view(ctx.input_0_shape)
+            )
+
+        # Store the current states.
+        bwd_cpu_rng_state = torch.get_rng_state()
+        bwd_cuda_rng_state = torch.cuda.get_rng_state()
+        bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()
+
+        # Set the states to what it used to be before the forward pass.
+        torch.set_rng_state(ctx.fwd_cpu_rng_state)
+        _set_cuda_rng_state(ctx.fwd_cuda_rng_state)
+        get_cuda_rng_tracker().set_states(ctx.fwd_cuda_rng_state_tracker)
+
+        # Compute the forward pass.
+        detached_inputs = detach_variable(inputs)
+        detached_ori_outputs = detach_variable(fa_output)
+        detached_ori_outputs[0].requires_grad = True #only 0 element need grad in output of FA: [Tensor output, Tensor logsumexp, Tensor dropout_mask]
+        # ori_outputs is not requires_grad
+        with torch.enable_grad():
+            outputs_before_fa = ctx.run_function(*detached_inputs)
+            # outputs_before_fa: query, key, value, attention_mask, {"attn_mask_type":attn_mask_type, "attention_bias":attention_bias, "packed_seq_params":packed_seq_params}
+            outputs_before_fa = ctx.last_function.forward_before_fa(*outputs_before_fa[:4], **outputs_before_fa[4])
+            outputs = ctx.last_function.forward_after_fa(detached_ori_outputs,
+                                                         ctx.qkv_format, ctx.indices_q,
+                                                         ctx.batch_size, ctx.attn_mask_type,
+                                                         ctx.max_seqlen_q, ctx.q_shape, ctx.v_shape)
+        # Set the states back to what it was at the start of this function.
+        torch.set_rng_state(bwd_cpu_rng_state)
+        _set_cuda_rng_state(bwd_cuda_rng_state)
+        get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker)
+
+
+        if isinstance(outputs, torch.Tensor):
+            outputs = (outputs,)
+        # filter out non tensor outputs for backward pass
+        outputs, args = zip(*filter(lambda x: torch.is_tensor(x[0]), zip(outputs, args)))
+        torch.autograd.backward(outputs, args)
+
+        #costum bwd fa
+        with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False):
+            with torch.no_grad():
+                grad_input = torch.ops.aten._scaled_dot_product_attention_flash_musa_backward(
+                    # ori_outputs[0][0].grad,
+                    detached_ori_outputs[0].grad,
+                    *outputs_before_fa[:3], #q, k, v
+                    *detached_ori_outputs, #(Tensor output, Tensor logsumexp, Tensor dropout_mask)
+                    is_causal="causal" in ctx.attn_mask_type, #causal same as fwd
+                )
+
+        #bwd before fa: for qkv
+        torch.autograd.backward(outputs_before_fa[:3], grad_input)
+        grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs)
+        return (None, None, None) + grads
+
+
+def checkpointViranceAttention(run_function, last_function, distribute_saved_activations, *args):
+    """Checkpoint a model or part of the model.
+    This has been directly copied from torch.utils.checkpoint."""
+    return CheckpointFunctionViranceAttention.apply(run_function, last_function, distribute_saved_activations, *args)
+# HACK(huang.huang)
+
+
+from transformer_engine.musa.pytorch.utils import add_attr
+from megatron.core import tensor_parallel
+add_attr(tensor_parallel, 'checkpointVirance', checkpointVirance)
+add_attr(tensor_parallel, 'checkpointViranceAttention', checkpointViranceAttention)
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/recomupte_variance/schedules.py b/megatron-lm-musa-patch/musa_patch/recomupte_variance/schedules.py
new file mode 100644
index 000000000..65c904e6b
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/recomupte_variance/schedules.py
@@ -0,0 +1,256 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+import contextlib
+from typing import Iterator, List, Union
+
+import torch
+
+from megatron.core import parallel_state
+from megatron.core.enums import ModelType
+from megatron.core.transformer.moe.router import MoEAuxLossAutoScaler
+from megatron.core.utils import (
+    get_attr_wrapped_model,
+    get_model_type,
+)
+from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
+from megatron.core.pipeline_parallel.schedules import set_current_microbatch, custom_backward
+
+# Types
+Shape = Union[List[int], torch.Size]
+
+
+#HACK(huang.huang): add FP8GlobalStateManager.reduce_and_update_fp8_tensors to the end of forward and backward,
+# to avoid redundant calls of reduce among dp
+def forward_step(
+    forward_step_func,
+    data_iterator,
+    model,
+    num_microbatches,
+    input_tensor,
+    forward_data_store,
+    config,
+    collect_non_loss_data=False,
+    checkpoint_activations_microbatch=None,
+    is_first_microbatch=False,
+    current_microbatch=None,
+    encoder_decoder_xattn=False,
+):
+    """Forward step for passed-in model.
+
+    If it is the first stage, the input tensor is obtained from the data_iterator.
+    Otherwise, the passed-in input_tensor is used.
+
+    Args:
+        forward_step_func (callable):
+            The forward step function for the model that takes the
+            data iterator as the first argument, and model as the second.
+            This user's forward step is expected to output a tuple of two elements:
+
+                1. The output object from the forward step. This output object needs to be a
+                    tensor or some kind of collection of tensors. The only hard requirement
+                    for this object is that it needs to be acceptible as input into the second
+                    function.
+                2. A function to reduce (optionally) the output from the forward step. This
+                    could be a reduction over the loss from the model, it could be a function that
+                    grabs the output from the model and reformats, it could be a function that just
+                    passes through the model output. This function must have one of the following
+                    patterns, and depending on the pattern different things happen internally:
+
+                        a. A tuple of reduced loss and some other data. Note that in this case
+                            the first argument is divided by the number of global microbatches,
+                            assuming it is a loss, so that the loss is stable as a function of
+                            the number of devices the step is split across.
+                        b. A triple of reduced loss, number of tokens, and some other data. This
+                            is similar to case (a), but the loss is further averaged across the
+                            number of tokens in the batch. If the user is not already averaging
+                            across the number of tokens, this pattern is useful to use.
+                        c. Any arbitrary data the user wants (eg a dictionary of tensors, a list
+                            of tensors, etc in the case of inference). To trigger case 3 you need
+                            to specify `collect_non_loss_data=True` and you may also want to
+                            specify `forward_only=True` in the call to the parent forward_backward
+                            function.
+        data_iterator (iterator):
+            The data iterator.
+        model (nn.Module):
+            The model to perform the forward step on.
+        num_microbatches (int):
+            The number of microbatches.
+        input_tensor (Tensor or list[Tensor]):
+            The input tensor(s) for the forward step.
+        forward_data_store (list):
+            The list to store the forward data. If you go down path 2.a or
+            2.b for the return of your forward reduction function then this will store only the
+            final dimension of the output, for example the metadata output by the loss function.
+            If you go down the path of 2.c then this will store the entire output of the forward
+            reduction function applied to the model output.
+        config (object):
+            The configuration object.
+        collect_non_loss_data (bool, optional):
+            Whether to collect non-loss data. Defaults to False.
+            This is the path to use if you want to collect arbitrary output from the model forward,
+            such as with inference use cases. Defaults to False.
+        checkpoint_activations_microbatch (int, optional):
+            The microbatch to checkpoint activations.
+            Defaults to None.
+        is_first_microbatch (bool, optional):
+            Whether it is the first microbatch. Defaults to False.
+        current_microbatch (int, optional):
+            The current microbatch. Defaults to None.
+
+    Returns:
+        Tensor or list[Tensor]: The output object(s) from the forward step.
+        Tensor: The number of tokens.
+    """
+    if config.timers is not None:
+        config.timers('forward-compute', log_level=2).start()
+
+    if is_first_microbatch and hasattr(model, 'set_is_first_microbatch'):
+        model.set_is_first_microbatch()
+    if current_microbatch is not None:
+        set_current_microbatch(model, current_microbatch)
+
+    unwrap_output_tensor = False
+    if not isinstance(input_tensor, list):
+        input_tensor = [input_tensor]
+        unwrap_output_tensor = True
+
+    set_input_tensor = get_attr_wrapped_model(model, "set_input_tensor")
+    set_input_tensor(input_tensor)
+
+    if config.enable_autocast:
+        context_manager = torch.autocast("cuda", dtype=config.autocast_dtype)
+    else:
+        context_manager = contextlib.nullcontext()
+    with context_manager:
+        if checkpoint_activations_microbatch is None:
+            output_tensor, loss_func = forward_step_func(data_iterator, model)
+        else:
+            output_tensor, loss_func = forward_step_func(
+                data_iterator, model, checkpoint_activations_microbatch
+            )
+
+    num_tokens = torch.tensor(0, dtype=torch.int)
+    if parallel_state.is_pipeline_last_stage():
+        if not collect_non_loss_data:
+            outputs = loss_func(output_tensor)
+            if len(outputs) == 3:
+                output_tensor, num_tokens, loss_reduced = outputs
+                if not config.calculate_per_token_loss:
+                    output_tensor /= num_tokens
+                    output_tensor /= num_microbatches
+            else:
+                # preserve legacy loss averaging behavior (ie, over the number of microbatches)
+                assert len(outputs) == 2
+                output_tensor, loss_reduced = outputs
+                output_tensor /= num_microbatches
+            forward_data_store.append(loss_reduced)
+        else:
+            data = loss_func(output_tensor, non_loss_data=True)
+            forward_data_store.append(data)
+    FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=True, skip=False)
+
+    if config.timers is not None:
+        config.timers('forward-compute').stop()
+
+    # Set the loss scale for the auxiliary loss of the MoE layer.
+    # Since we use a trick to do backward on the auxiliary loss, we need to set the scale
+    # explicitly.
+    if hasattr(config, 'num_moe_experts') and config.num_moe_experts is not None:
+        # Calculate the loss scale based on the grad_scale_func if available, else default to 1.
+        loss_scale = (
+            config.grad_scale_func(torch.ones(1, device=output_tensor.device))
+            if config.grad_scale_func is not None
+            else torch.tensor(1.0)
+        )
+        # Set the loss scale
+        MoEAuxLossAutoScaler.set_loss_scale(loss_scale / num_microbatches)
+
+    # If T5 model and in decoder stack, then send encoder_hidden_state
+    # downstream as well.
+    model_type = get_model_type(model)
+    if (
+        model_type == ModelType.encoder_and_decoder
+        and encoder_decoder_xattn
+        and parallel_state.is_inside_decoder()
+    ):
+        return [output_tensor, input_tensor[-1]], num_tokens
+
+    if unwrap_output_tensor:
+        return output_tensor, num_tokens
+    return [output_tensor], num_tokens
+
+
+def backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config):
+    """Backward step through passed-in output tensor.
+
+    If last stage, output_tensor_grad is None, otherwise gradient of loss
+    with respect to stage's output tensor.
+
+    Returns gradient of loss with respect to input tensor (None if first
+    stage)."""
+
+    # NOTE: This code currently can handle at most one skip connection. It
+    # needs to be modified slightly to support arbitrary numbers of skip
+    # connections.
+
+    if config.timers is not None:
+        config.timers('backward-compute', log_level=2).start()
+
+    # Retain the grad on the input_tensor.
+    unwrap_input_tensor_grad = False
+    if not isinstance(input_tensor, list):
+        input_tensor = [input_tensor]
+        unwrap_input_tensor_grad = True
+    for x in input_tensor:
+        if x is not None:
+            x.retain_grad()
+
+    if not isinstance(output_tensor, list):
+        output_tensor = [output_tensor]
+    if not isinstance(output_tensor_grad, list):
+        output_tensor_grad = [output_tensor_grad]
+
+    # Backward pass.
+    if output_tensor_grad[0] is None and config.grad_scale_func is not None:
+        output_tensor[0] = config.grad_scale_func(output_tensor[0])
+
+    if config.deallocate_pipeline_outputs:
+        custom_backward(output_tensor[0], output_tensor_grad[0])
+    else:
+        torch.autograd.backward(output_tensor[0], grad_tensors=output_tensor_grad[0])
+
+    # Collect the grad of the input_tensor.
+    input_tensor_grad = [None]
+    if input_tensor is not None:
+        input_tensor_grad = []
+        for x in input_tensor:
+            if x is None:
+                input_tensor_grad.append(None)
+            else:
+                input_tensor_grad.append(x.grad)
+
+    # Handle single skip connection if it exists (encoder_hidden_state in
+    # model with encoder and decoder).
+    if (
+        parallel_state.get_pipeline_model_parallel_world_size() > 1
+        and model_type == ModelType.encoder_and_decoder
+        and len(output_tensor_grad) > 1  # excludes models that lack a skip connection.
+    ):
+        if output_tensor_grad[1] is not None:
+            assert input_tensor_grad[-1] is not None
+            input_tensor_grad[-1].add_(output_tensor_grad[1])
+    if unwrap_input_tensor_grad:
+        input_tensor_grad = input_tensor_grad[0]
+
+    FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False, skip=False)
+    if config.timers is not None:
+        config.timers('backward-compute').stop()
+
+    return input_tensor_grad
+#HACK(huang.huang)
+
+
+from transformer_engine.musa.pytorch.utils import replace_attr
+from megatron.core.pipeline_parallel import schedules
+replace_attr(schedules, "forward_step", forward_step)
+replace_attr(schedules, "backward_step", backward_step)
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/recomupte_variance/transformer_engine.py b/megatron-lm-musa-patch/musa_patch/recomupte_variance/transformer_engine.py
new file mode 100644
index 000000000..2d2fd5195
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/recomupte_variance/transformer_engine.py
@@ -0,0 +1,245 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import dataclasses
+import os
+from typing import Any, Callable, Optional
+
+import torch
+from packaging.version import Version as PkgVersion
+from torch import Tensor
+
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.parallel_state import (
+    get_context_parallel_global_ranks,
+    get_context_parallel_group,
+    get_hierarchical_context_parallel_groups,
+    get_tensor_model_parallel_group,
+)
+from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import get_te_version, is_te_min_version
+
+from megatron.core.extensions.transformer_engine import TEDotProductAttention
+
+
+# HACK(huang.huang): recompute-variance for fa: modify __init__ for TEDotProductAttention,
+# just add a attr "recompute_variance" when call super.__init__,
+# add functions "forward_after_fa", "forward_before_fa"
+def TEDotProductAttention__init__(
+    self,
+    config: TransformerConfig,
+    layer_number: int,
+    attn_mask_type: AttnMaskType,
+    attention_type: str,
+    attention_dropout: Optional[float] = None,
+    softmax_scale: Optional[float] = None,
+    k_channels: Optional[int] = None,
+    v_channels: Optional[int] = None,
+    cp_comm_type: str = "p2p",
+):
+    self.config = config
+    self.te_forward_mask_type = False
+    self.qkv_format: str = 'sbhd'
+
+    if self.config.apply_query_key_layer_scaling != bool(
+        int(os.getenv('NVTE_APPLY_QK_LAYER_SCALING', '0'))
+    ):
+        raise ValueError(
+            f"apply_query_key_layer_scaling is {self.config.apply_query_key_layer_scaling} "
+            f"but environment variable NVTE_APPLY_QK_LAYER_SCALING is "
+            f"{os.getenv('NVTE_APPLY_QK_LAYER_SCALING')}. Transformer Engine does not support "
+            f"setting query key layer scaling via argument, so these two must match."
+        )
+
+    extra_kwargs: dict[str, Any] = {}
+    if is_te_min_version("0.11.0"):
+        extra_kwargs["num_gqa_groups"] = self.config.num_query_groups
+    elif self.config.num_query_groups != self.config.num_attention_heads:
+        raise ValueError(
+            f"Transformer Engine v{get_te_version()} does not support Grouped Query Attention, "
+            f"use a newer version of Transformer Engine. "
+            f"(num_query_groups ({self.config.num_query_groups}) != "
+            f"num_attention_heads ({self.config.num_attention_heads}))"
+        )
+
+    if is_te_min_version("0.10.0"):
+        extra_kwargs["attention_type"] = attention_type
+        # older version don't need attention_type
+
+    if is_te_min_version("0.12.0", check_equality=False):
+        self.te_forward_mask_type = True
+
+    # This check is important as CP config can be disabled while having a valid CP group
+    # Example - Disabling CP for encoder while a valid CP group exists for decoder
+    if self.config.context_parallel_size > 1:
+        assert is_te_min_version(
+            "1.0.0"
+        ), "Only Transformer-Engine version >= 1.0.0 supports context parallelism!"
+        if getattr(TEDotProductAttention, "cp_stream") is None:
+            TEDotProductAttention.cp_stream = torch.cuda.Stream()
+        extra_kwargs["cp_group"] = get_context_parallel_group(check_initialized=False)
+        extra_kwargs["cp_global_ranks"] = get_context_parallel_global_ranks(
+            check_initialized=False
+        )
+        extra_kwargs["cp_stream"] = TEDotProductAttention.cp_stream
+        if is_te_min_version("1.10.0"):
+            if cp_comm_type is None:
+                extra_kwargs["cp_comm_type"] = "p2p"
+            elif cp_comm_type == "a2a+p2p":
+                assert is_te_min_version("1.12.0"), (
+                    f"Transformer-Engine v{get_te_version()} must be >= 1.12.0 to support"
+                    "hierarchical cp commucation."
+                )
+                extra_kwargs["cp_comm_type"] = "a2a+p2p"
+                extra_kwargs["cp_group"] = get_hierarchical_context_parallel_groups(
+                    check_initialized=False
+                )
+            else:
+                extra_kwargs["cp_comm_type"] = cp_comm_type
+
+    if self.config.deterministic_mode:
+        if int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")) != 0:
+            raise RuntimeError(
+                "deterministic_mode is on and we are using DotProductAttention from "
+                "Transformer Engine, but NVTE_ALLOW_NONDETERMINISTIC_ALGO is not 0. "
+                f"Currently set to: {os.getenv('NVTE_ALLOW_NONDETERMINISTIC_ALGO', 'not set')}."
+            )
+
+    if config.window_size is not None:
+        # Check version
+        assert is_te_min_version("1.2.0"), (
+            f"Transformer-Engine v{get_te_version()} must be >= 1.2.0 to support"
+            "sliding window attention."
+        )
+        extra_kwargs['window_size'] = config.window_size
+
+    if is_te_min_version("1.10.0"):
+        # TE 1.10.0 introduces the ability to set the different k and v channels
+        kv_channels = (
+            (k_channels, v_channels)
+            if k_channels is not None and v_channels is not None
+            else self.config.kv_channels
+        )
+        extra_kwargs['softmax_scale'] = softmax_scale
+    else:
+        kv_channels = self.config.kv_channels
+
+    self.kept_packed_seq_params = set(
+        field.name for field in dataclasses.fields(PackedSeqParams)
+    )
+    if get_te_version() < PkgVersion("1.3.0"):
+        # TE 1.3.0 introduces precomputing max_seqlen to remove unnecessary kernels and D2H
+        # copies (#555)
+        # These two arguments did not exist prior to 1.3.0
+        self.kept_packed_seq_params.discard("max_seqlen_q")
+        self.kept_packed_seq_params.discard("max_seqlen_kv")
+
+    if get_te_version() < PkgVersion("1.10.0"):
+        # TE 1.8.0 introduces cu_seqlens_padded which is the cu_seqlens with paddings counted
+        # in each individual sequence in THD format dataset
+        # These two arguments did not exist prior to 1.8.0. Full support added in 1.10.0 (#1012)
+        self.kept_packed_seq_params.discard("cu_seqlens_q_padded")
+        self.kept_packed_seq_params.discard("cu_seqlens_kv_padded")
+
+    super(TEDotProductAttention, self).__init__(
+        num_attention_heads=self.config.num_attention_heads,
+        kv_channels=kv_channels,
+        attention_dropout=(
+            self.config.attention_dropout if attention_dropout is None else attention_dropout
+        ),
+        attn_mask_type=attn_mask_type.name,
+        sequence_parallel=self.config.sequence_parallel,
+        tp_size=self.config.tensor_model_parallel_size,
+        get_rng_state_tracker=(
+            get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+        ),
+        tp_group=get_tensor_model_parallel_group(check_initialized=False),
+        layer_number=layer_number,
+        recompute_variance = self.config.recompute_variance, # MUSA patch: support recompute_variance
+        **extra_kwargs,
+    )
+
+def TEDotProductAttention_forward_before_fa(self,
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    attention_mask: Tensor,
+    attn_mask_type: AttnMaskType,
+    attention_bias: Tensor = None,
+    packed_seq_params: PackedSeqParams = None,):
+
+    packed_seq_kwargs = (
+        {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
+        if packed_seq_params is not None
+        else {}
+    )
+    # overwrite self.qkv_format depending on self.config.apply_rope_fusion, which can be set
+    # after init
+    if self.config.apply_rope_fusion and is_te_min_version("0.13.0", check_equality=False):
+        self.qkv_format = 'bshd'
+
+    qkv_format = packed_seq_kwargs.get('qkv_format', self.qkv_format)
+
+    # WAR for peak memory usage.
+    # See https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/2388
+    if self.config.apply_rope_fusion and qkv_format == 'bshd':
+        query, key, value = [x.transpose(0, 1).contiguous() for x in (query, key, value)]
+        # In PyTorch, the following two tensors are in fact the same:
+        #   Tensor with shape (1, S, H, D) and stride (S*H*D, H*D, D, 1)
+        #   Tensor with shape (1, S, H, D) and stride (H*D, H*D, D, 1)
+        # Stride for a dimension that is 1 has no meaning, so tensors created two different ways
+        # can have same shape but different strides.
+        # We unify them to the first one to pass the stride check in TE
+        if value.shape == key.shape and value.shape[0] == 1 and value.stride() != key.stride():
+            value = value.as_strided(value.shape, key.stride())
+
+    attention_bias_kwargs = {}
+    if attention_bias is not None:
+        assert is_te_min_version("1.2.0"), (
+            f"Transformer-Engine v{get_te_version()} must be >= 1.2.0 to support"
+            "`attention_bias`."
+        )
+        attention_bias_kwargs = dict(
+            core_attention_bias_type='post_scale_bias', core_attention_bias=attention_bias
+        )
+    if self.te_forward_mask_type:
+        if qkv_format == 'thd' and is_te_min_version("1.7.0"):
+            # thd format uses flash attention with cuDNN kernel which requires is_padding=True,
+            # so the only acceptable mask types are `padding_causal` and `padding`. These do not
+            # necessarily indicate there are padded tokens in the sequence.
+            if attn_mask_type == AttnMaskType.causal:
+                attn_mask_type = AttnMaskType.padding_causal
+            elif attn_mask_type == AttnMaskType.no_mask:
+                attn_mask_type = AttnMaskType.padding
+        output = super(TEDotProductAttention, self).forward_before_fa(
+            query,
+            key,
+            value,
+            attention_mask,
+            attn_mask_type=attn_mask_type.name,
+            **attention_bias_kwargs,
+            **packed_seq_kwargs,
+        )
+    else:
+        output = super(TEDotProductAttention, self).forward_before_fa(
+            query, key, value, attention_mask, **attention_bias_kwargs, **packed_seq_kwargs
+        )
+
+    return output
+
+def TEDotProductAttention_forward_after_fa(self, *args):
+    core_attn_out = super(TEDotProductAttention, self).forward_after_fa(*args)
+
+    if self.config.apply_rope_fusion and self.qkv_format == 'bshd':
+        return core_attn_out.transpose(0, 1)
+    else:
+        return core_attn_out
+# HACK(huang.huang)
+
+
+from transformer_engine.musa.pytorch.utils import replace_attr, add_attr
+
+replace_attr(TEDotProductAttention,"__init__", TEDotProductAttention__init__)
+add_attr(TEDotProductAttention, "forward_before_fa", TEDotProductAttention_forward_before_fa)
+add_attr(TEDotProductAttention, "forward_after_fa", TEDotProductAttention_forward_after_fa)
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/recomupte_variance/transformer_layer.py b/megatron-lm-musa-patch/musa_patch/recomupte_variance/transformer_layer.py
new file mode 100644
index 000000000..bd2735282
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/recomupte_variance/transformer_layer.py
@@ -0,0 +1,192 @@
+from megatron.core.utils import make_viewless_tensor
+from megatron.core import parallel_state, tensor_parallel
+
+from transformer_engine.pytorch.distributed import checkpoint, checkpointVirance
+
+# HACK(huang.huang): support mlp_rms_recompute and mla_rms_recompute,
+# which need to decide to do layernorm in TransformerLayer or inner mlp/mla
+def TransformerLayer_forward(
+    self,
+    hidden_states,
+    attention_mask=None,
+    context=None,
+    context_mask=None,
+    rotary_pos_emb=None,
+    rotary_pos_cos=None,
+    rotary_pos_sin=None,
+    attention_bias=None,
+    inference_params=None,
+    packed_seq_params=None,
+    sequence_len_offset=None,
+):
+    """
+    Perform a forward pass through the transformer layer.
+
+    This method implements the core computation of a transformer layer, including
+    self-attention, cross-attention (if applicable), and feed-forward operations.
+
+    Args:
+        hidden_states (Tensor): Input tensor of shape [s, b, h] where s is sequence length,
+            b is batch size, and h is hidden size.
+        attention_mask (Tensor): Mask tensor for self-attention.
+        context (Tensor, optional): Context tensor for cross-attention.
+        context_mask (Tensor, optional): Mask tensor for cross-attention.
+        rotary_pos_emb (Tensor, optional): Rotary positional embeddings.
+        attention_bias (Tensor, optional): Bias tensor for Q * K.T.
+        inference_params (object, optional): Parameters for inference-time optimizations.
+        packed_seq_params (object, optional): Parameters for packed sequence processing.
+
+    Returns:
+        Tuple[Tensor, Tensor]: A tuple containing:
+            output (Tensor): Transformed hidden states of shape [s, b, h].
+            context (Tensor): Updated context tensor if cross-attention is used,
+            otherwise None.
+    """
+
+    # Residual connection.
+    residual = hidden_states
+
+    # Optional Input Layer norm
+    #HACK(huang.haung): support mla_rms_recompute
+    if self.config.mla_rms_recompute:
+        assert self.config.attn_recompute, 'mla_rms_recompute only use with attn_recompute now.'
+        def rms_with_down_proj(hidden_states):
+            hidden_states = self.input_layernorm(hidden_states)
+            if self.self_attention.config.q_lora_rank is not None:
+                q_compressed, _ = self.self_attention.linear_q_down_proj(hidden_states)
+            else:
+                q_compressed = hidden_states
+            kv_combined, _ = self.self_attention.linear_kv_down_proj(hidden_states)
+            return q_compressed, kv_combined
+        input_layernorm_output = None
+        if self.config.fp8:
+            if self.config.recompute_variance == True:
+                linears = (self.self_attention.linear_q_down_proj, self.self_attention.linear_kv_down_proj)
+                q_compressed, kv_combined = checkpointVirance(
+                    self.input_layernorm,
+                    linears,
+                    hidden_states,
+                    distribute_saved_activations=self.config.distribute_saved_activations,
+                    get_rng_state_tracker=tensor_parallel.random.get_cuda_rng_tracker,
+                    tp_group=parallel_state.get_tensor_model_parallel_group(),
+                )
+            else:
+                q_compressed, kv_combined =  checkpoint(
+                    rms_with_down_proj,
+                    hidden_states,
+                    distribute_saved_activations=self.config.distribute_saved_activations,
+                    get_rng_state_tracker=tensor_parallel.random.get_cuda_rng_tracker,
+                    tp_group=parallel_state.get_tensor_model_parallel_group(),
+                )
+        else:
+            if self.config.recompute_variance:
+                assert self.self_attention.config.q_lora_rank is not None, "not support Now" #TODO
+                linears = (self.self_attention.linear_q_down_proj, self.self_attention.linear_kv_down_proj)
+                q_compressed, kv_combined = tensor_parallel.checkpointVirance(
+                    self.input_layernorm,
+                    linears,
+                    False,
+                    hidden_states)
+            else:
+                q_compressed, kv_combined =  tensor_parallel.checkpoint(
+                    rms_with_down_proj, False, hidden_states)
+
+        attention_output_with_bias = self.self_attention(
+            input_layernorm_output,
+            attention_mask=attention_mask,
+            inference_params=inference_params,
+            rotary_pos_emb=rotary_pos_emb,
+            rotary_pos_cos=rotary_pos_cos,
+            rotary_pos_sin=rotary_pos_sin,
+            attention_bias=attention_bias,
+            packed_seq_params=packed_seq_params,
+            sequence_len_offset=sequence_len_offset,
+            q_compressed=q_compressed,
+            kv_combined=kv_combined,
+        )
+
+    else: #maintain original implement, to support non MLA attention
+        input_layernorm_output = self.input_layernorm(hidden_states)
+        # Self attention.
+        attention_output_with_bias = self.self_attention(
+            input_layernorm_output,
+            attention_mask=attention_mask,
+            inference_params=inference_params,
+            rotary_pos_emb=rotary_pos_emb,
+            rotary_pos_cos=rotary_pos_cos,
+            rotary_pos_sin=rotary_pos_sin,
+            attention_bias=attention_bias,
+            packed_seq_params=packed_seq_params,
+            sequence_len_offset=sequence_len_offset,
+        )
+    ## HACK(huang.haung)
+
+    # TODO: could we move `bias_dropout_add_exec_handler` itself
+    # inside the module provided in the `bias_dropout_add_spec` module?
+    with self.bias_dropout_add_exec_handler():
+        hidden_states = self.self_attn_bda(self.training, self.config.bias_dropout_fusion)(
+            attention_output_with_bias, residual, self.hidden_dropout
+        )
+
+    # Residual connection.
+    residual = hidden_states
+
+    # Optional Layer norm after self-attention
+    pre_cross_attn_layernorm_output = self.pre_cross_attn_layernorm(hidden_states)
+
+    # Cross attention.
+    attention_output_with_bias = self.cross_attention(
+        pre_cross_attn_layernorm_output,
+        attention_mask=context_mask,
+        key_value_states=context,
+        inference_params=inference_params,
+    )
+
+    if isinstance(attention_output_with_bias, dict) and "context" in attention_output_with_bias:
+        context = attention_output_with_bias["context"]
+
+    # TODO: could we move `bias_dropout_add_exec_handler` itself
+    # inside the module provided in the `bias_dropout_add_spec` module?
+    with self.bias_dropout_add_exec_handler():
+        hidden_states = self.cross_attn_bda(self.training, self.config.bias_dropout_fusion)(
+            attention_output_with_bias, residual, self.hidden_dropout
+        )
+
+    # Residual connection.
+    residual = hidden_states
+
+    # Optional Layer norm post the cross-attention.
+    #HACK(huang.haung): support mlp_rms_recompute
+    if self.config.mlp_rms_recompute:
+        pre_mlp_layernorm_output = None
+        mlp_output_with_bias = self.mlp(hidden_states, self.pre_mlp_layernorm)
+    else:
+        pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states)
+        mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
+    ## HACK(huang.haung)
+    # TODO: could we move `bias_dropout_add_exec_handler` itself
+    # inside the module provided in the `bias_dropout_add_spec` module?
+    with self.bias_dropout_add_exec_handler():
+        hidden_states = self.mlp_bda(self.training, self.config.bias_dropout_fusion)(
+            mlp_output_with_bias, residual, self.hidden_dropout
+        )
+
+    # Jit compiled function creates 'view' tensor. This tensor
+    # potentially gets saved in the MPU checkpoint function context,
+    # which rejects view tensors. While making a viewless tensor here
+    # won't result in memory savings (like the data loader, or
+    # p2p_communication), it serves to document the origin of this
+    # 'view' tensor.
+    output = make_viewless_tensor(
+        inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True
+    )
+
+    # CUDA graph requires returned values to be Tensors
+    if self.config.external_cuda_graph and self.training:
+        return output
+    return output, context
+## HACK(huang.huang)
+
+from transformer_engine.musa.pytorch.utils import replace_attr
+from megatron.core.transformer.transformer_layer import TransformerLayer
+replace_attr(TransformerLayer, "forward", TransformerLayer_forward)
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/rotary_pos_embedding.py b/megatron-lm-musa-patch/musa_patch/rotary_pos_embedding.py
new file mode 100644
index 000000000..24681a02e
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/rotary_pos_embedding.py
@@ -0,0 +1,142 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from megatron.core.transformer.transformer_config import TransformerConfig
+    from megatron.core.transformer.transformer_block import TransformerBlock
+
+import logging
+
+import torch
+import torch_musa
+from torch import Tensor, nn
+
+from megatron.core import parallel_state
+
+logger = logging.getLogger(__name__)
+
+try:
+    from apex.transformer.functional import (
+        fused_apply_rotary_pos_emb,
+        fused_apply_rotary_pos_emb_thd,
+    )
+
+    HAVE_APPLY_ROPE_FUSION = True
+except ImportError:
+    HAVE_APPLY_ROPE_FUSION = False
+
+def _rotate_half(x: Tensor, rotary_interleaved: bool) -> Tensor:
+    """Change sign so the last dimension becomes [-odd, +even]
+
+    Args:
+        x (Tensor): Input tensor
+
+    Returns:
+        Tensor: Tensor rotated half
+    """
+    if not rotary_interleaved:
+        x1, x2 = torch.chunk(x, 2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+    else:
+        x1 = x[:, :, :, ::2]
+        x2 = x[:, :, :, 1::2]
+        x_new = torch.stack((-x2, x1), dim=-1)
+        return x_new.view(x_new.shape[0], x_new.shape[1], x_new.shape[2], -1)
+
+def apply_rotary_pos_emb_bshd(t: Tensor, freqs: Tensor, rotary_interleaved: bool = False) -> Tensor:
+    """Apply rotary positional embedding to input tensor T.
+
+    check https://kexue.fm/archives/8265 for detailed formulas
+
+    Args:
+        t (Tensor): Input tensor T is of shape [seq_length, ... , dim]
+        freqs (Tensor): Rotary Positional embedding tensor freq is of shape [seq_length, ..., dim]
+
+    Returns:
+        Tensor: The input tensor after applying RoPE
+    """
+    rot_dim = freqs.shape[-1]
+
+    # ideally t_pass is empty so rotary pos embedding is applied to all tensor t
+    t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
+
+    # first part is cosine component
+    # second part is sine component, need to change signs with _rotate_half method
+    cos_ = torch.cos(freqs).to(t.dtype)
+    sin_ = torch.sin(freqs).to(t.dtype)
+
+    t = (t * cos_) + (_rotate_half(t, rotary_interleaved) * sin_)
+    return torch.cat((t, t_pass), dim=-1)
+
+
+def apply_rotary_pos_emb_thd(
+    t: Tensor, cu_seqlens: Tensor, freqs: Tensor, rotary_interleaved: bool = False
+) -> Tensor:
+    """A baseline implementation of applying RoPE for `thd` format.
+
+    Args:
+        t (Tensor): Input tensor T is of shape [t, h, d]
+        cu_seqlens(Tensor):  Cumulative sum of sequence lengths in a batch for `t`,
+        with shape [b + 1] and dtype torch.int32.
+        freqs (Tensor): Rotary Positional embedding tensor freq is of shape [max_s, 1, 1, d]
+
+    Returns:
+        Tensor: Shape [t, h, d]. The input tensor after applying RoPE.
+    """
+
+    seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+    return torch.cat(
+        [
+            apply_rotary_pos_emb_bshd(x.unsqueeze(1), freqs[: x.size(0)])
+            for x in torch.split(t, seqlens)
+        ]
+    ).squeeze(1)
+
+def apply_rotary_pos_emb(
+    t: Tensor, freqs: Tensor, config: TransformerConfig, cu_seqlens: Optional[Tensor] = None, mscale: float = 1.0,
+):
+
+    """
+    Reroute to the appropriate apply_rotary_pos_emb function depending on
+    fused/unfused kernels, or bshd (conventional) / thd (packed seq) format
+    """
+    # assert cu_seqlens is None, "Only support cu_seqlens is None for now!"
+    if config.apply_rope_fusion and not HAVE_APPLY_ROPE_FUSION:
+        # setting apply_rope_fusion in config to False so that subsequent queries to this config also return False
+        config.apply_rope_fusion = False
+        if not getattr(apply_rotary_pos_emb, "printed_fused_warning", False):
+            logger.warning(
+                "Setting apply_rope_fusion to false because its implementation"
+                " is not included in Apex. Try upgrading to the latest version"
+            )
+            apply_rotary_pos_emb.printed_fused_warning = True
+    if config.apply_rope_fusion:
+        if cu_seqlens is None:
+            if config.multi_latent_attention:
+                return torch.rope(t, freqs.squeeze(1).squeeze(1), rotary_interleaved=False, batch_first=False, multi_latent_attention=True)
+            else:
+                return fused_apply_rotary_pos_emb(t, freqs, transpose_output_memory=True)
+        else:
+            return fused_apply_rotary_pos_emb_thd(t, cu_seqlens, freqs)
+    else:
+        if cu_seqlens is None:
+            return apply_rotary_pos_emb_bshd(t, freqs, rotary_interleaved=config.rotary_interleaved)
+        else:
+            return apply_rotary_pos_emb_thd(
+                t, cu_seqlens, freqs, rotary_interleaved=config.rotary_interleaved
+            )
+
+# import megatron.core.models.common.embeddings.rotary_pos_embedding
+# megatron.core.models.common.embeddings.rotary_pos_embedding.apply_rotary_pos_emb = apply_rotary_pos_emb
+# import megatron.core.transformer.attention
+# megatron.core.transformer.attention.apply_rotary_pos_emb = apply_rotary_pos_emb
+
+import sys
+for k in sys.modules:
+    if k.startswith('megatron.core'):
+        for target in ['apply_rotary_pos_emb']:
+            if getattr(sys.modules[k], target, None):
+                setattr(sys.modules[k], target, apply_rotary_pos_emb)
diff --git a/megatron-lm-musa-patch/musa_patch/router.py b/megatron-lm-musa-patch/musa_patch/router.py
new file mode 100644
index 000000000..1bc897495
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/router.py
@@ -0,0 +1,55 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from functools import partial
+import torch
+
+from .moe_utils import (
+    sequence_load_balancing_loss_func,
+    topk_softmax_with_capacity,
+)
+
+def seq_aux_loss_load_balancing(self, logits: torch.Tensor, bsz: int, seq_length: int):
+    """Apply loss-based load balancing to the logits tensor."""
+
+    probs, routing_map, tokens_per_expert = topk_softmax_with_capacity(
+            logits,
+            self.topk,
+            capacity_factor=self.config.moe_expert_capacity_factor,
+            pad_to_capacity=self.config.moe_pad_expert_input_to_capacity,
+            drop_policy=self.config.moe_token_drop_policy,
+            use_pre_softmax=self.config.moe_router_pre_softmax,
+            num_groups=self.config.moe_router_num_groups,
+            group_topk=self.config.moe_router_group_topk,
+            scaling_factor=self.config.moe_router_topk_scaling_factor,
+            deterministic_mode=self.config.deterministic_mode,
+            score_function=self.score_function,
+            expert_bias=self.expert_bias,
+            device_level_capacity=self.config.moe_device_level_capacity,
+        )
+
+    if self.training:
+        if self.score_function == "sigmoid":
+            scores = torch.sigmoid(logits)
+        else:
+            scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
+        aux_loss_func = partial(
+            sequence_load_balancing_loss_func,
+            probs=scores,
+            routing_map=routing_map,
+            batch_size=bsz,
+            seq_length=seq_length,
+            topk=self.topk,
+            moe_router_topk_limited_devices=self.config.moe_router_group_topk,
+            moe_device_level_aux_loss_coeff=self.config.moe_device_level_aux_loss_coeff,
+            moe_comm_aux_loss_coeff=self.config.moe_comm_aux_loss_coeff,
+            moe_complementary_seq_aux_loss=self.config.moe_complementary_seq_aux_loss,
+        )
+        probs = self.apply_load_balancing_loss(
+            activation=probs, load_balancing_loss_func=aux_loss_func
+        )
+
+    return probs, routing_map
+
+
+import megatron.core.transformer.moe.router
+megatron.core.transformer.moe.router.TopKRouter.seq_aux_loss_load_balancing = seq_aux_loss_load_balancing
diff --git a/megatron-lm-musa-patch/musa_patch/theoretical_memory_usage.py b/megatron-lm-musa-patch/musa_patch/theoretical_memory_usage.py
new file mode 100644
index 000000000..2f977bd18
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/theoretical_memory_usage.py
@@ -0,0 +1,220 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+"""Computes theoretical memory footprint for model training."""
+
+
+import math
+
+NUM_BYTES_IN_MEGABYTE = 1024 * 1024
+
+def compute_weight_and_optimizer_memory(args, verbose=False):
+    # Attention projection size.
+    attn_dim = args.kv_channels
+    kv_projection_size = attn_dim * args.num_attention_heads
+    query_projection_size = attn_dim * args.num_attention_heads
+    ## MLA
+    if args.kv_lora_rank:
+        kv_projection_size = attn_dim * args.num_attention_heads + args.kv_lora_rank
+    if args.q_lora_rank:
+        query_projection_size = attn_dim * args.num_attention_heads + args.q_lora_rank
+
+    output_projection_size = attn_dim * args.num_attention_heads
+    ## Group Query Attention.
+    if args.group_query_attention:
+        kv_projection_size = args.num_query_groups / args.num_attention_heads * kv_projection_size
+    else:
+        attn_size = 1
+    attn_size = query_projection_size + 2 * kv_projection_size + output_projection_size
+    attn_multiplier = attn_size / 2 / args.hidden_size
+
+    # swiglu
+    gated_linear_multiplier = 3 / 2 if args.swiglu else 1
+
+    # MoE or Dense
+    num_experts = 0 if args.num_experts is None else args.num_experts
+    shared_expert_ffn_hidden_size = 0 if args.moe_shared_expert_intermediate_size else args.moe_shared_expert_intermediate_size
+    num_shared_expert = args.moe_shared_expert_intermediate_size // args.moe_ffn_hidden_size
+    num_moe_layer = 0 if args.moe_layer_freq is None else len(args.moe_layer_freq)
+    num_dense_layer = args.num_layers - num_moe_layer
+    mlp_multiplier_dense = num_dense_layer * args.ffn_hidden_size
+    mlp_multiplier_moe =  num_moe_layer * num_experts * args.moe_ffn_hidden_size
+    mlp_multiplier_shared_expert = num_moe_layer * num_shared_expert * shared_expert_ffn_hidden_size
+    mlp_multiplier = (mlp_multiplier_dense + mlp_multiplier_moe + mlp_multiplier_shared_expert) /args.num_layers/args.hidden_size
+
+    num_parameters_in_transformer_layers = (
+        2
+        * args.num_layers
+        * args.hidden_size
+        * args.hidden_size
+        * (
+            # Attention.
+            attn_multiplier
+            # MLP.
+            # + ((args.ffn_hidden_size / args.hidden_size) * num_experts * gated_linear_multiplier)
+            + mlp_multiplier * gated_linear_multiplier
+            # Router
+            +  num_experts / args.hidden_size * 2
+            # Transformer layernorms.
+            + (2 / args.hidden_size)
+            # Final layernorm.
+            + (1 / (args.num_layers * args.hidden_size))
+        )
+    )
+
+    embedding_size = args.hidden_size * args.padded_vocab_size
+    if args.untie_embeddings_and_output_weights:
+        num_parameters_in_embedding_layers = 2 * embedding_size
+    else:
+        num_parameters_in_embedding_layers = embedding_size
+
+    # TODO; add MTP block and projection
+    num_parameters_in_mtp = 0
+
+    num_total_parameters = num_parameters_in_transformer_layers + num_parameters_in_embedding_layers + num_parameters_in_mtp
+    if verbose:
+        print(
+            f"Number of parameters in transformer layers in billions: "
+            f"{num_parameters_in_transformer_layers / 10**9: .2f}"
+        )
+        print(
+            f"Number of parameters in embedding layers in billions: "
+            f"{num_parameters_in_embedding_layers / 10**9:.2f}"
+        )
+        print(f"Total number of parameters in billions: {num_total_parameters / 10**9:.2f}")
+
+    # Most loaded model shard has (1/pp_size transformer layers + 1 embedding layer) / tp_size.
+    num_parameters_on_most_loaded_model_shard = (
+        (num_parameters_in_transformer_layers / args.pipeline_model_parallel_size) + embedding_size
+    ) / args.tensor_model_parallel_size
+    if args.untie_embeddings_and_output_weights and args.pipeline_model_parallel_size == 1:
+        num_parameters_on_most_loaded_model_shard += (
+            embedding_size / args.tensor_model_parallel_size
+        )
+    if verbose:
+        print(
+            f"Number of parameters in most loaded shard in billions: "
+            f"{num_parameters_on_most_loaded_model_shard / 10**9:.4f}"
+        )
+
+    if args.pipeline_model_parallel_size > 1:
+        # Other shards just have (1/pp_size transformer layers) / tp_size.
+        num_parameters_on_other_model_shards = num_parameters_in_transformer_layers / (
+            args.pipeline_model_parallel_size * args.tensor_model_parallel_size
+        )
+        if verbose:
+            print(
+                f"Number of parameters in other shards in billions: "
+                f"{num_parameters_on_other_model_shards / 10**9:.4f}"
+            )
+
+    num_bytes_per_parameter = (
+        18 if not args.use_distributed_optimizer else 6 + (12 / args.data_parallel_size)
+    )
+    weight_and_optimizer_memory = (
+        num_parameters_on_most_loaded_model_shard * num_bytes_per_parameter
+    )
+
+    return weight_and_optimizer_memory
+
+
+def compute_activation_memory(args, num_microbatches, verbose=False):
+    # Using formula in Table 2 of https://arxiv.org/pdf/2205.05198.pdf.
+    # We are trying to compute the maximum activation footprint, so all calculations in this
+    # function are for the first pipeline stage.
+
+    # TODO: This function needs to take into account query_projection_size potentially being
+    # different from hidden_size.
+
+    # Memory footprint from transformer layer (self-attention and MLP).
+    activation_memory = (args.seq_length * args.micro_batch_size * args.hidden_size) * (
+        18 + (4 * (args.ffn_hidden_size / args.hidden_size))
+    )
+    if verbose:
+        print(
+            f"Activation memory footprint per transformer layer: "
+            f"{activation_memory / NUM_BYTES_IN_MEGABYTE / args.tensor_model_parallel_size:.1f} MB"
+        )
+    activation_memory *= args.num_layers
+
+    # Now add activation memory required for input embeddings, last LayerNorm and output layer.
+
+    # Input to embedding (pp_size microbatches in flight).
+    activation_memory += (
+        8 * args.seq_length * args.micro_batch_size * args.pipeline_model_parallel_size
+    )
+    # Dropout in embedding layer (pp_size microbatches in flight).
+    activation_memory += (
+        args.seq_length
+        * args.micro_batch_size
+        * args.hidden_size
+        * args.pipeline_model_parallel_size
+    )
+
+    # Multiply by interleaved PP memory factor.
+    if args.virtual_pipeline_model_parallel_size is not None:
+        interleaved_schedule_memory_penalty = 1 + (
+            (args.pipeline_model_parallel_size - 1)
+            / (args.pipeline_model_parallel_size * args.virtual_pipeline_model_parallel_size)
+        )
+        in_flight_microbatches = math.ceil(
+            interleaved_schedule_memory_penalty * args.pipeline_model_parallel_size
+        )
+        if verbose:
+            print(
+                f"Memory penalty from interleaved schedule: {interleaved_schedule_memory_penalty:.2f}"
+            )
+            print(f"Number of in-flight microbatches: {in_flight_microbatches}")
+        activation_memory *= interleaved_schedule_memory_penalty
+
+    # If using non-interleaved schedule, number of microbatches in pipeline can be less than pp_size,
+    # so discount accordingly.
+    if args.virtual_pipeline_model_parallel_size is None and args.pipeline_model_parallel_size > 1:
+        if num_microbatches is not None:
+            activation_memory *= min(1, num_microbatches / args.pipeline_model_parallel_size)
+            in_flight_microbatches = min(num_microbatches, args.pipeline_model_parallel_size)
+        else:
+            in_flight_microbatches = args.pipeline_model_parallel_size
+        if verbose:
+            print(f"Number of in-flight microbatches: {in_flight_microbatches}")
+
+    if args.pipeline_model_parallel_size == 1:
+        # Inputs to output layer and CE loss.
+        activation_memory += (
+            args.seq_length
+            * args.micro_batch_size
+            * args.hidden_size
+            * 4
+            * (1 + (args.padded_vocab_size / args.hidden_size))
+        )
+
+    # Activation memory is partitioned by TP size due to tensor and sequence model parallelism.
+    return activation_memory / args.tensor_model_parallel_size
+
+
+def report_theoretical_memory(args, num_microbatches=None, verbose=False):
+    weight_and_optimizer_memory = (
+        compute_weight_and_optimizer_memory(args, verbose=verbose) / NUM_BYTES_IN_MEGABYTE
+    )
+
+    # Formulae here assume sequence parallelism and selective activation recomputation.
+    if not args.sequence_parallel or args.recompute_granularity != 'selective':
+        print(
+            f"Theoretical memory footprints: weight and optimizer={weight_and_optimizer_memory:.2f} MB"
+        )
+        return
+
+    activation_memory = (
+        compute_activation_memory(args, num_microbatches=num_microbatches, verbose=verbose)
+        / NUM_BYTES_IN_MEGABYTE
+    )
+    total_memory = weight_and_optimizer_memory + activation_memory
+
+    print(
+        f"Theoretical memory footprints: weight and optimizer={weight_and_optimizer_memory:.2f} MB, "
+        f"activation={activation_memory:.2f} MB, total={total_memory:.2f} MB\n"
+    )
+
+import megatron.training.theoretical_memory_usage
+megatron.training.theoretical_memory_usage.compute_weight_and_optimizer_memory = compute_weight_and_optimizer_memory
+megatron.training.theoretical_memory_usage.report_theoretical_memory = report_theoretical_memory
+megatron.training.theoretical_memory_usage.compute_activation_memory = compute_activation_memory
diff --git a/megatron-lm-musa-patch/musa_patch/token_dispatcher.py b/megatron-lm-musa-patch/musa_patch/token_dispatcher.py
new file mode 100644
index 000000000..2ef68f737
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/token_dispatcher.py
@@ -0,0 +1,462 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from typing import List, Optional, Tuple
+
+import torch
+
+from megatron.core.tensor_parallel import (
+    all_to_all,
+    gather_from_sequence_parallel_region,
+    reduce_scatter_to_sequence_parallel_region,
+)
+from megatron.core.transformer.moe.moe_utils import (
+    get_capacity,
+    permute,
+    sort_chunks_by_idxs,
+    unpermute,
+)
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+#HACK(huang.huang):use pin_memory, non_blocking, cuda.stream to optimze D2H
+#class MoEAlltoAllTokenDispatcher:
+#   function update: __init__, preprocess, token_permutation, token_unpermutation
+#       __init__: extra opration after ori_init
+#           a. copy sort_input_by_local_experts and restore_output_by_local_experts to cpu, prevent redunt D2H in each Forward
+#           b. add self.cuda_dtoh_stream to support async copy
+#           c. follow Megatron:main, add cuda_sync_point_priority and cuda_dtoh_point to manage sync point
+#       preprocess:
+#           a. remove D2H such as .cpu() and do it in token_permutation as once
+#           b. use new method to manage sync point
+#           c. don't set self.cuda_sync_point to "before_permutation_1" when not used drop and pad, while it's seems to be unnecessary
+#       token_permutation:
+#           a. do ravel for num_global_tokens_per_local_expert and num_global_tokens_per_local_expert_T before D2H
+#           b. pass numpy to sort_chunks instead call tolist for gpu tensor
+#       token_unpermutation:
+#           a. pass numpy to sort_chunks instead call tolist for gpu tensor
+#
+#   new function: _maybe_update_cuda_sync_point, _ensure_cpu_tensor, _maybe_dtoh_and_synchronize
+#       _maybe_update_cuda_sync_point: update sync_point if this point is reached earlier than current one
+#       _ensure_cpu_tensor:
+#           a. use pin_memory to alloc cpu tensor, and copy value in gpu tensor to it
+#           b. will be used under cuda_dtoh_stream to prevent block
+#       _maybe_dtoh_and_synchronize:
+#           a. call _ensure_cpu_tensor in dtoh_point, which equals to "before_pemutation1" expert_capacity used
+#           b. in self.cuda_sync_point, do cuda_dtoh_stream.synchronize(), .numpy, tolist where synchronize needs to be done before
+#           .numpy and tolist, otherwise D2H maybe not be finished.
+def MoEAlltoAllTokenDispatcher___init__(
+    self, num_local_experts: int, local_expert_indices: List[int], config: TransformerConfig
+) -> None:
+    self._orig___init__(num_local_experts, local_expert_indices, config)
+
+    self.sort_input_by_local_experts = self.sort_input_by_local_experts.cpu().numpy()
+    self.restore_output_by_local_experts = self.restore_output_by_local_experts.cpu().numpy()
+    self.cuda_sync_point_priority = {
+        "before_permutation_1": 0,
+        "before_ep_alltoall": 1,
+        "before_permutation_2": 2,
+        "before_finish": 3,
+        "no_sync": 4,
+    }
+    self.cuda_dtoh_point = "before_permutation_1"
+    self.cuda_dtoh_stream = torch.cuda.Stream()
+
+def MoEAlltoAllTokenDispatcher__maybe_update_cuda_sync_point(self, point: str):
+    """
+    Update the CUDA sync point if the priority of the new point is higher than the current
+    sync point, which means the new point is reached earlier than the current sync point.
+    """
+    if (
+        self.cuda_sync_point_priority[point]
+        < self.cuda_sync_point_priority[self.cuda_sync_point]
+    ):
+        self.cuda_sync_point = point
+
+def MoEAlltoAllTokenDispatcher__ensure_cpu_tensor(self, cpu_attr_name, gpu_tensor):
+    if gpu_tensor is None:
+        return
+    cpu_tensor = getattr(self, cpu_attr_name, None)
+    if cpu_tensor is None:
+        cpu_tensor = torch.empty(
+            gpu_tensor.size(),
+            device="cpu",
+            pin_memory=True,
+            dtype=gpu_tensor.dtype
+        )
+        setattr(self, cpu_attr_name, cpu_tensor)
+    cpu_tensor.copy_(gpu_tensor, non_blocking=True)
+    gpu_tensor.record_stream(torch.cuda.current_stream())
+
+def MoEAlltoAllTokenDispatcher__maybe_dtoh_and_synchronize(
+    self, point: str, tokens_per_expert: torch.Tensor = None,
+    num_global_tokens_per_local_expert: torch.Tensor = None,
+    num_global_tokens_per_local_expert_T: torch.Tensor = None,
+):
+    """
+    Move all possible GPU tensors to CPU and make a synchronization at the expected point.
+    """
+    if not self.drop_and_pad:
+        if point == self.cuda_dtoh_point:
+            # Move all possible GPU tensors to CPU at self.cuda_dtoh_point.
+            on_side_stream = torch.cuda.current_stream() != self.cuda_dtoh_stream
+            if on_side_stream:
+                self.cuda_dtoh_stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(self.cuda_dtoh_stream):
+                # TODO: use MemcpyBatchAsync instead.
+                self._ensure_cpu_tensor('tokens_per_expert_cpu', tokens_per_expert)
+                if self.ep_size > 1 or self.tp_size > 1:
+                    self._ensure_cpu_tensor('output_splits_tp_cpu', self.output_splits_tp)
+                    self._ensure_cpu_tensor('input_splits_cpu', self.input_splits)
+                    self._ensure_cpu_tensor('output_splits_cpu', self.output_splits)
+                #NOTE(huang.huang): only fuse-moe-permute need num_out_tokens, but in that case fused-kernel will do D2H itself,
+                # if we want to do this D2H here, then we need to sync stream before permute1, which resulted in D2H not being well overlaped
+                # self._ensure_cpu_tensor('num_out_tokens_cpu', self.num_out_tokens)
+                if self.num_local_experts > 1:
+                    self._ensure_cpu_tensor('num_global_tokens_per_local_expert_cpu', num_global_tokens_per_local_expert)
+                    self._ensure_cpu_tensor('num_global_tokens_per_local_expert_T_cpu', num_global_tokens_per_local_expert_T)
+
+        if point == self.cuda_sync_point:
+            # Synchronize with the dtoh stream at self.cuda_sync_point.
+            self.cuda_dtoh_stream.synchronize()
+            # Need to do before sync, otherwise copy for value in num_global_tokens_per_local_expert_cpu is not finished
+            # self.num_out_tokens = self.num_out_tokens_cpu.numpy()
+
+            tokens_per_expert = self.tokens_per_expert_cpu.numpy().copy()
+            # need copy(), because recompute groupedLinear1 save it as msplit, value in ctx will change while next copy gpu data to _cpu
+            # not use tolist, since expert will call tolist again, otherwise we need to modify experts.py
+
+            if self.ep_size > 1 or self.tp_size > 1:
+                self.output_splits_tp = self.output_splits_tp_cpu.numpy().tolist()
+                self.input_splits = self.input_splits_cpu.numpy().tolist()
+                self.output_splits = self.output_splits_cpu.numpy().tolist()
+            if self.num_local_experts > 1:
+                self.num_global_tokens_per_local_expert = self.num_global_tokens_per_local_expert_cpu.numpy()
+                self.num_global_tokens_per_local_expert_T = self.num_global_tokens_per_local_expert_T_cpu.numpy()
+
+    return tokens_per_expert
+
+def MoEAlltoAllTokenDispatcher_preprocess(self, routing_map: torch.Tensor) -> torch.Tensor:
+    """
+    Preprocess token routing map for AlltoAll communication and token permutation.
+
+    This method computes the number of tokens assigned to each expert based on the routing_map.
+    It also initializes the necessary data structures for AlltoAll communication, such as input
+    and output splits, and the mapping between global tokens and local experts.
+
+    Args:
+        routing_map (torch.Tensor): The mapping of tokens to experts, with shape
+            [num_tokens, num_experts].
+
+    Returns:
+        torch.Tensor: Tensor containing the number of tokens assigned to local expert.
+    """
+    # [num_experts], number of tokens assigned to each expert from the current rank's input.
+    num_local_tokens_per_expert = routing_map.sum(dim=0).long()
+
+    if self.drop_and_pad:
+        # Drop and pad the input to capacity.
+        num_tokens = routing_map.size(0) * self.config.moe_router_topk
+        self.capacity = get_capacity(
+            num_tokens=num_tokens,
+            num_experts=self.num_experts,
+            capacity_factor=self.moe_expert_capacity_factor,
+        )
+        self.num_out_tokens = self.capacity * self.num_experts
+        # [num_local_experts], number of tokens processed by each expert.
+        num_tokens_per_local_expert = torch.full(
+            (self.num_local_experts,),
+            self.capacity * self.tp_size * self.ep_size,
+            dtype=torch.long,
+        )
+        # [tp_size * ep_size, num_local_experts]. Represents the number of tokens sent
+        # to each local expert by all ranks.
+        self.num_global_tokens_per_local_expert = torch.full(
+            (self.num_experts * self.tp_size,),
+            self.capacity,
+            dtype=torch.long,
+            device=self.permute_idx_device,
+        )
+        return num_tokens_per_local_expert
+    elif self.config.moe_expert_capacity_factor is not None:
+        # Drop tokens to capacity, no padding.
+        # A synchronization is needed before the first
+        # permutation to get the `num_out_tokens` CPU value.
+        self.num_out_tokens = num_local_tokens_per_expert.sum()
+        # #TODO(huang.huang): make sure num_out_tokens is not needed for permutation_1 excpet drop_and_pad
+        # self.cuda_sync_point = "before_permutation_1"
+
+    else:
+        # Dropless
+        self.num_out_tokens = routing_map.size(0) * self.config.moe_router_topk
+        #HACK(huang.huang): move setattr for self.cuda_sync_point below
+        # if self.ep_size > 1 or self.num_local_experts > 1:
+        #     # Token dropless and enable ep. A synchronization is needed before expert parallel
+        #     # AlltoAll communication to get the `input_splits` and `output_splits` CPU values.
+        #     self.cuda_sync_point = "before_ep_alltoall"
+        # else:
+        #     # Token dropless and no ep. A synchronization is needed before the returns
+        #     # to get the `tokens_per_expert` CPU value for
+        #     self.cuda_sync_point = "before_finish"
+        ##HACK(huang.huang)
+    if self.ep_size > 1 or self.tp_size > 1:
+        # ===================================================
+        # Calculate input_splits, output_splits for alltoall/allgather in variable size.
+        # ===================================================
+        # [ep_size]. Represents the number of tokens sent by the current rank to other
+        # EP ranks.
+        self.input_splits = num_local_tokens_per_expert.reshape(
+            self.ep_size, self.num_local_experts
+        ).sum(axis=1)
+        # Gather the global distribution of tokens across ranks.
+        # num_global_tokens_per_expert represents the number of tokens sent to each
+        # expert by all ranks.
+        # [tp_size, ep_size, num_experts]
+        num_global_tokens_per_expert = (
+            gather_from_sequence_parallel_region(
+                num_local_tokens_per_expert, group=self.tp_ep_group
+            )
+            .reshape(self.ep_size, self.tp_size, self.num_experts)
+            .transpose(0, 1)
+        )
+        # [tp_size, ep_size, num_experts] -> [tp_size, ep_size, num_local_experts]
+        num_global_tokens_per_local_expert = num_global_tokens_per_expert[
+            :, :, self.local_expert_indices[0] : self.local_expert_indices[-1] + 1
+        ].contiguous()
+        # [tp_size, ep_size, num_local_experts] -> [tp_size, ep_size]
+        num_global_tokens_per_rank = num_global_tokens_per_local_expert.sum(axis=2)
+        # [tp_size, ep_size] -> [ep_size]
+        # self.output_splits represents the number of tokens received by the current rank
+        # from other EP rank.
+        self.output_splits = num_global_tokens_per_rank[self.tp_rank]
+        # [tp_size, ep_size] -> [tp_size]
+        # self.output_splits_tp represents the number of tokens received by the current
+        # rank from other TP rank.
+        self.output_splits_tp = num_global_tokens_per_rank.sum(axis=1)
+        # [tp_size, ep_size, num_local_experts] -> [num_local_experts]
+        num_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=(0, 1))
+        # A synchronization is needed before expert parallel AlltoAll communication
+        # to get the `input_splits` and `output_splits` CPU values.
+        self._maybe_update_cuda_sync_point("before_ep_alltoall")
+    else:
+        num_global_tokens_per_local_expert = num_local_tokens_per_expert.reshape(
+            self.num_experts
+        )
+        num_tokens_per_local_expert = num_local_tokens_per_expert
+        # A synchronization is needed before the returns
+        # to get the `num_tokens_per_local_expert` CPU value.
+        self._maybe_update_cuda_sync_point("before_finish")
+
+    if self.num_local_experts > 1:
+        # [tp_size * ep_size, num_local_experts]. Represents the number of tokens sent
+        # to each local expert by all ranks.
+        self.num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view(
+            -1, self.num_local_experts
+        )
+        if not self.config.moe_permute_fusion:
+            # A synchronization is needed before permutation 2
+            # to get the `num_global_tokens_per_local_expert` CPU value.
+            self._maybe_update_cuda_sync_point("before_permutation_2")
+
+    return num_tokens_per_local_expert
+
+
+def MoEAlltoAllTokenDispatcher_token_permutation(
+    self, hidden_states: torch.Tensor, probs: torch.Tensor, routing_map: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Dispatch tokens to local experts using AlltoAll communication.
+
+    This method performs the following steps:
+    1. Preprocess the routing map to get metadata for communication and permutation.
+    2. Permute input tokens for AlltoAll communication.
+    3. Perform expert parallel AlltoAll communication.
+    4. Sort tokens by local expert (if multiple local experts exist).
+
+    Args:
+        hidden_states (torch.Tensor): Input token embeddings.
+        probs (torch.Tensor): The probabilities of token to experts assignment.
+        routing_map (torch.Tensor): The mapping of token to experts assignment.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]:
+            - Permuted token embeddings for local experts.
+            - Number of tokens per expert.
+    """
+    # Preprocess: Get the metadata for communication, permutation and computation operations.
+    self.hidden_shape = hidden_states.shape
+    self.probs = probs
+    self.routing_map = routing_map
+    assert probs.dim() == 2, "Expected 2D tensor for probs"
+    assert routing_map.dim() == 2, "Expected 2D tensor for token2expert mask"
+    assert routing_map.dtype == torch.bool, "Expected bool tensor for mask"
+    hidden_states = hidden_states.view(-1, self.hidden_shape[-1])
+    tokens_per_expert = self.preprocess(self.routing_map)
+
+    if self.shared_experts is not None:
+        self.shared_experts.pre_forward_comm(hidden_states.view(self.hidden_shape))
+
+    # Permutation 1: input to AlltoAll input
+    self.hidden_shape_before_permute = hidden_states.shape
+
+    #GPU operations that data depended on need to be performed before the d2h command
+    num_global_tokens_per_local_expert = self.num_global_tokens_per_local_expert.ravel()
+    num_global_tokens_per_local_expert_T = self.num_global_tokens_per_local_expert.T.ravel()
+
+    # Permutation 1: input to AlltoAll input
+    tokens_per_expert = self._maybe_dtoh_and_synchronize(
+        "before_permutation_1", tokens_per_expert,
+        num_global_tokens_per_local_expert, num_global_tokens_per_local_expert_T
+    )
+    permutated_local_input_tokens, self.reversed_local_input_permutation_mapping = permute(
+        hidden_states,
+        routing_map,
+        num_out_tokens=self.num_out_tokens,
+        fused=self.config.moe_permute_fusion,
+        drop_and_pad=self.drop_and_pad,
+    )
+    # Perform expert parallel AlltoAll communication
+    tokens_per_expert = self._maybe_dtoh_and_synchronize(
+        "before_ep_alltoall", tokens_per_expert,
+        num_global_tokens_per_local_expert, num_global_tokens_per_local_expert_T
+    )
+
+    global_input_tokens = all_to_all(
+        self.ep_group, permutated_local_input_tokens, self.output_splits, self.input_splits
+    )
+    if self.shared_experts is not None:
+        self.shared_experts.linear_fc1_forward_and_act(global_input_tokens)
+
+    if self.tp_size > 1:
+        if self.output_splits_tp is None:
+            output_split_sizes = None
+        else:
+            output_split_sizes = self.output_splits_tp.tolist()
+        global_input_tokens = gather_from_sequence_parallel_region(
+            global_input_tokens, group=self.tp_group, output_split_sizes=output_split_sizes
+        )
+
+    # Permutation 2: Sort tokens by local expert.
+    tokens_per_expert = self._maybe_dtoh_and_synchronize(
+        "before_permutation_2", tokens_per_expert,
+        num_global_tokens_per_local_expert, num_global_tokens_per_local_expert_T
+    )
+    if self.num_local_experts > 1:
+        if self.drop_and_pad:
+            global_input_tokens = (
+                global_input_tokens.view(
+                    self.tp_size * self.ep_size,
+                    self.num_local_experts,
+                    self.capacity,
+                    *global_input_tokens.size()[1:],
+                )
+                .transpose(0, 1)
+                .contiguous()
+                .flatten(start_dim=0, end_dim=2)
+            )
+        else:
+            global_input_tokens = sort_chunks_by_idxs(
+                global_input_tokens,
+                self.num_global_tokens_per_local_expert,
+                self.sort_input_by_local_experts,
+                fused=self.config.moe_permute_fusion,
+            )
+
+    tokens_per_expert = self._maybe_dtoh_and_synchronize(
+        "before_finish", tokens_per_expert,
+        num_global_tokens_per_local_expert, num_global_tokens_per_local_expert_T
+    )
+    return global_input_tokens, tokens_per_expert
+
+
+def MoEAlltoAllTokenDispatcher_token_unpermutation(
+    self, hidden_states: torch.Tensor, bias: Optional[torch.Tensor] = None
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """
+    Reverse the token permutation to restore the original order.
+
+    This method performs the following steps:
+    1. Unsort tokens by local expert (if multiple local experts exist).
+    2. Perform expert parallel AlltoAll communication to restore the original order.
+    3. Unpermute tokens to restore the original order.
+
+    Args:
+        hidden_states (torch.Tensor): Output from local experts.
+        bias (torch.Tensor, optional): Bias tensor (not supported).
+
+    Returns:
+        Tuple[torch.Tensor, Optional[torch.Tensor]]:
+            - Unpermuted token embeddings in the original order.
+            - None (bias is not supported).
+    """
+    assert bias is None, "Bias is not supported in MoEAlltoAllTokenDispatcher"
+
+    # Unpermutation 2: Unsort tokens by local expert.
+    if self.num_local_experts > 1:
+        if self.drop_and_pad:
+            hidden_states = (
+                hidden_states.view(
+                    self.num_local_experts,
+                    self.tp_size * self.ep_size,
+                    self.capacity,
+                    *hidden_states.size()[1:],
+                )
+                .transpose(0, 1)
+                .contiguous()
+                .flatten(start_dim=0, end_dim=2)
+            )
+        else:
+            hidden_states = sort_chunks_by_idxs(
+                hidden_states,
+                self.num_global_tokens_per_local_expert_T,
+                self.restore_output_by_local_experts,
+                fused=self.config.moe_permute_fusion,
+            )
+
+    if self.tp_size > 1:
+        if self.output_splits_tp is None:
+            input_split_sizes = None
+        else:
+            input_split_sizes = self.output_splits_tp.tolist()
+        hidden_states = reduce_scatter_to_sequence_parallel_region(
+            hidden_states, group=self.tp_group, input_split_sizes=input_split_sizes
+        )
+
+    # Perform expert parallel AlltoAll communication
+    # hidden_states: [SEQL, H] -> [SEQL, H/TP]
+    permutated_local_input_tokens = all_to_all(
+        self.ep_group, hidden_states, self.input_splits, self.output_splits
+    )
+    if self.shared_experts is not None:
+        self.shared_experts.linear_fc2_forward(permutated_local_input_tokens)
+        self.shared_experts.post_forward_comm()
+
+    # Unpermutation 1: AlltoAll output to output
+    output = unpermute(
+        permutated_local_input_tokens,
+        self.reversed_local_input_permutation_mapping,
+        restore_shape=self.hidden_shape_before_permute,
+        probs=self.probs,
+        routing_map=self.routing_map,
+        fused=self.config.moe_permute_fusion,
+        drop_and_pad=self.drop_and_pad,
+    )
+
+    # Reshape the output tensor
+    output = output.view(self.hidden_shape)
+
+    # Add shared experts output
+    if self.shared_experts is not None:
+        shared_expert_output = self.shared_experts.get_output()
+        output += shared_expert_output
+    return output, None
+##HACK(huang.huang)
+
+from transformer_engine.musa.pytorch.utils import replace_attr, add_attr
+from megatron.core.transformer.moe.token_dispatcher import MoEAlltoAllTokenDispatcher
+replace_attr(MoEAlltoAllTokenDispatcher, '__init__', MoEAlltoAllTokenDispatcher___init__)
+add_attr(MoEAlltoAllTokenDispatcher, '_maybe_update_cuda_sync_point', MoEAlltoAllTokenDispatcher__maybe_update_cuda_sync_point)
+add_attr(MoEAlltoAllTokenDispatcher, '_ensure_cpu_tensor', MoEAlltoAllTokenDispatcher__ensure_cpu_tensor)
+add_attr(MoEAlltoAllTokenDispatcher, '_maybe_dtoh_and_synchronize', MoEAlltoAllTokenDispatcher__maybe_dtoh_and_synchronize)
+replace_attr(MoEAlltoAllTokenDispatcher, 'preprocess', MoEAlltoAllTokenDispatcher_preprocess)
+replace_attr(MoEAlltoAllTokenDispatcher, 'token_permutation', MoEAlltoAllTokenDispatcher_token_permutation)
+replace_attr(MoEAlltoAllTokenDispatcher, 'token_unpermutation', MoEAlltoAllTokenDispatcher_token_unpermutation)
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/training.py b/megatron-lm-musa-patch/musa_patch/training.py
new file mode 100644
index 000000000..ed740ce38
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/training.py
@@ -0,0 +1,954 @@
+from datetime import datetime
+
+import gc
+import os
+import sys
+import logging
+import torch
+import torch.distributed
+from megatron.core import mpu
+
+from megatron.core.transformer.moe.moe_utils import track_moe_metrics
+from megatron.core.transformer.moe.router import TopKRouter
+from megatron.training.global_vars import (
+    get_args,
+    get_timers,
+    get_tensorboard_writer,
+    get_wandb_writer,
+    get_one_logger
+)
+# get_num_microbatches
+from megatron.core.num_microbatches_calculator import get_num_microbatches
+from megatron.training.utils import (
+    report_memory,
+    print_rank_last
+)
+from megatron.core.rerun_state_machine import (
+    get_rerun_state_machine,
+)
+from megatron.core.utils import (
+    check_param_hashes_across_dp_replicas,
+)
+from megatron.training.theoretical_memory_usage import report_theoretical_memory
+from megatron.training import one_logger_utils
+from megatron.training.initialize import write_args_to_tensorboard
+from megatron.core.distributed import finalize_model_grads
+from megatron.core.distributed import DistributedDataParallel as DDP
+
+from megatron.training.training import (
+    print_datetime,
+    save_checkpoint_and_time,
+    train_step,
+    evaluate_and_print_results,
+)
+from megatron.training.async_utils import maybe_finalize_async_save
+from megatron.core.num_microbatches_calculator import (
+    get_current_global_batch_size,
+    get_current_running_global_batch_size,
+    get_num_microbatches,
+    update_num_microbatches
+)
+from megatron.training.utils import (
+    calc_params_l2_norm,
+    print_rank_0,
+    print_rank_last,
+    report_memory,
+)
+from megatron.training import ft_integration
+from megatron.training.global_vars import (
+    get_args,
+    get_timers,
+    get_tensorboard_writer,
+    get_wandb_writer,
+    get_one_logger
+)
+from .profiling import maybe_enable_profiling
+
+from megatron.training.training import (
+    enable_forward_pre_hook,
+    disable_forward_pre_hook,
+    post_training_step_callbacks,
+    checkpoint_and_decide_exit
+)
+
+from megatron.training.utils import (
+    calc_params_l2_norm,
+    logical_and_across_model_parallel_group,
+    reduce_max_stat_across_model_parallel_group,
+    print_rank_0,
+    print_rank_last,
+    report_memory,
+    unwrap_model,
+)
+
+from megatron.core.pipeline_parallel import get_forward_backward_func
+try:
+    import mlflow
+except Exception as e:
+    print(f"import mlflow failed {str(e)}")
+
+logger = logging.getLogger(__name__)
+
+def throughput_calculator(args, elapsed_time_per_iter, consumed_tokens_per_iter):
+    # training_time = elapsed_time
+    system_throughput = float(consumed_tokens_per_iter) / elapsed_time_per_iter
+    world_size = args.world_size
+    chip_throughput = system_throughput / world_size
+    # For 70B
+    # all_param_num = getattr(args, "all_param_num", None)
+    # assert all_param_num is not None, "please set all_param_num"
+    # MFU = chip_throughput * 6 * all_param_num * (1 + args.seq_length / (6 * args.hidden_size) ) / 98e12
+    # # tflops_throughput = chip_throughput / float(config.flops_16bit) * 1e12
+    # # logger.info("Throughput(token per chip per second): " + str(chip_throughput))
+    # # logger.info("MFU: " + str(MFU))
+    # # logger.info("Throughput(token per TFLOPS): " + str(tflops_throughput))
+    h = args.hidden_size
+    s = args.seq_length
+    N = 12 * args.num_layers * h **2
+    D = 1
+
+    attn_matmul = 2 * N * D
+    attn_sdp = N * D * (s / h)
+    mlp_matmul = 4 * N * D
+    # moe
+    if args.num_experts is None:
+        factor = 1
+    else:
+        factor = args.moe_router_topk
+    activated_dense_flops = attn_matmul + attn_sdp + mlp_matmul * factor
+    if args.num_experts is not None:
+        act_params = N + args.num_layers *(args.num_experts - 1) * 8 * h**2
+        if torch.distributed.get_rank() == 0:
+            print(f"N: {N} Act param: {act_params} Act flops: {activated_dense_flops}")
+    tflops =  chip_throughput *  activated_dense_flops
+    mfu = tflops / 98e12
+
+    return chip_throughput, mfu
+
+def num_floating_point_operations(args, batch_size):
+    # Attention projection size.
+    query_projection_size = args.kv_channels * args.num_attention_heads
+    query_projection_to_hidden_size_ratio = query_projection_size / args.hidden_size
+    # Group Query Attention.
+    if not args.group_query_attention:
+        args.num_query_groups = args.num_attention_heads
+    # MoE.
+    num_experts_routed_to = 1 if args.num_experts is None else args.moe_router_topk
+    gated_linear_multiplier = 3 / 2 if args.swiglu else 1
+    shared_expert_ffn_hidden_size = (
+        0
+        if args.moe_shared_expert_intermediate_size is None
+        else args.moe_shared_expert_intermediate_size
+    )
+    if not args.multi_latent_attention:
+        return (
+            12
+            * batch_size
+            * args.seq_length
+            * args.num_layers
+            * args.hidden_size
+            * args.hidden_size
+            * (
+                # Attention.
+                (
+                    (
+                        1
+                        + (args.num_query_groups / args.num_attention_heads)
+                        + (args.seq_length / args.hidden_size)
+                    ) * query_projection_to_hidden_size_ratio
+                )
+                # MLP.
+                + (
+                    (args.moe_ffn_hidden_size / args.hidden_size)
+                    * num_experts_routed_to
+                    * gated_linear_multiplier
+                )
+                # Shared Experts.
+                + ((shared_expert_ffn_hidden_size / args.hidden_size) * gated_linear_multiplier)
+                # Logit.
+                + (args.padded_vocab_size / (2 * args.num_layers * args.hidden_size))
+            )
+        )
+    else:
+        if args.q_lora_rank is None:
+            mla_flops_q = args.hidden_size * args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+        else:
+            mla_flops_q = args.hidden_size * args.q_lora_rank +\
+                  args.q_lora_rank * args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+        return (
+            6
+            * batch_size
+            * args.seq_length
+            * args.num_layers
+            * (
+                # MLA Attention.
+                (
+                    (
+                        mla_flops_q
+                        + args.hidden_size * (args.kv_lora_rank + args.qk_pos_emb_head_dim)
+                        + args.num_attention_heads * args.kv_lora_rank * (args.qk_head_dim + args.v_head_dim)
+                        + args.num_attention_heads * args.seq_length * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+                        + args.num_attention_heads * args.seq_length * args.v_head_dim
+                        + args.num_attention_heads * args.v_head_dim * args.hidden_size
+                    )
+                )
+                # Router
+                + args.hidden_size * args.num_experts
+                # MLP.
+                + (
+                    2 * args.hidden_size *  args.moe_ffn_hidden_size * num_experts_routed_to
+                    * gated_linear_multiplier
+                )
+                # Shared Experts.
+                + (2 * args.hidden_size * shared_expert_ffn_hidden_size * gated_linear_multiplier)
+                # Logit.
+                + (args.padded_vocab_size * args.hidden_size / args.num_layers)
+            )
+        )
+
+def need_mlflow():
+    return os.getenv("MLFLOW_TRACKING_URI", default=None) and \
+            torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1)
+
+
+def train_step(forward_step_func, data_iterator,
+               model, optimizer, opt_param_scheduler, config):
+    """Single training step."""
+    args = get_args()
+    timers = get_timers()
+
+    rerun_state_machine = get_rerun_state_machine()
+    while rerun_state_machine.should_run_forward_backward(data_iterator):
+        # Set grad to zero.
+        for model_chunk in model:
+            model_chunk.zero_grad_buffer()
+        optimizer.zero_grad()
+
+        # Forward pass.
+        forward_backward_func = get_forward_backward_func()
+        losses_reduced = forward_backward_func( # forward_data_store
+            forward_step_func=forward_step_func,
+            data_iterator=data_iterator,
+            model=model,
+            num_microbatches=get_num_microbatches(),
+            seq_length=args.seq_length,
+            micro_batch_size=args.micro_batch_size,
+            decoder_seq_length=args.decoder_seq_length,
+            forward_only=False)
+    should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit()
+    if should_exit:
+        return {}, True, should_checkpoint, should_exit, exit_code, None, None
+
+    # Empty unused memory.
+    if args.empty_unused_memory_level >= 1:
+        torch.cuda.empty_cache()
+
+    # Vision gradients.
+    if args.vision_pretraining and args.vision_pretraining_type == "dino":
+        unwrapped_model = unwrap_model(model[0])
+        unwrapped_model.cancel_gradients_last_layer(args.curr_iteration)
+
+    # Update parameters.
+
+    timers('optimizer', log_level=1).start(barrier=args.barrier_with_L1_time)
+    update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
+    timers('optimizer').stop()
+
+    # when freezing sub-models we may have a mixture of successful and unsucessful ranks,
+    # so we must gather across mp ranks
+    update_successful = logical_and_across_model_parallel_group(update_successful)
+    # grad_norm and num_zeros_in_grad will be None on ranks without trainable params,
+    # so we must gather across mp ranks
+    grad_norm = reduce_max_stat_across_model_parallel_group(grad_norm)
+    if args.log_num_zeros_in_grad:
+        num_zeros_in_grad = reduce_max_stat_across_model_parallel_group(num_zeros_in_grad)
+
+    # Vision momentum.
+    if args.vision_pretraining and args.vision_pretraining_type == "dino":
+        unwrapped_model = unwrap_model(model[0])
+        unwrapped_model.update_momentum(args.curr_iteration)
+
+    # Update learning rate.
+    if update_successful:
+        increment = get_num_microbatches() * \
+                    args.micro_batch_size * \
+                    args.data_parallel_size
+        opt_param_scheduler.step(increment=increment)
+        skipped_iter = 0
+    else:
+        skipped_iter = 1
+
+    # Empty unused memory.
+    if args.empty_unused_memory_level >= 2:
+        torch.cuda.empty_cache()
+
+    if mpu.is_pipeline_last_stage(ignore_virtual=True):
+        # Average loss across microbatches.
+        loss_reduced = {}
+        for key in losses_reduced[0].keys():
+            numerator = 0
+            denominator = 0
+
+            # HACK(xuerong.huang): Reduce the report loss(loss_reduced) on the last training step of multi-microbatches.
+            if int(os.getenv("NO_LOSS_REDUCE", 0)):
+                val0 = losses_reduced[0][key]
+                if isinstance(val0, tuple) or isinstance(val0, list):
+                    reduce_data = [sum([v[key][0] for v in losses_reduced])] # get the sum of the losses of all microbatches
+                    reduce_data.extend([v[key][1] for v in losses_reduced]) # get the token-num of all microbatches
+                    reduce_data = torch.stack(reduce_data)
+                    torch.distributed.all_reduce(reduce_data, group=mpu.get_data_parallel_group()) # reduce the losses-sum and token-num from all dp-ranks
+                    numerator = reduce_data[0]
+                    denominator = sum(reduce_data[1:])
+                else:
+                    numerator = sum([v[key] for v in losses_reduced])
+                    denominator = len(losses_reduced)
+                    torch.distributed.all_reduce(numerator, group=mpu.get_data_parallel_group())    # reduce the losses-sum from all dp-ranks
+            # HACK(xuerong.huang): Reduce the report loss(loss_reduced) on the last training step of multi-microbatches.
+            else:
+                for x in losses_reduced:
+                    val = x[key]
+                    # there is one dict per microbatch. in new reporting, we average
+                    # over the total number of tokens across the global batch.
+                    if isinstance(val, tuple) or isinstance(val, list):
+                        numerator += val[0]
+                        denominator += val[1]
+                    else:
+                        # legacy behavior. we average over the number of microbatches,
+                        # and so the denominator is 1.
+                        numerator += val
+                        denominator += 1
+
+            loss_reduced[key] = numerator / denominator
+
+        return loss_reduced, skipped_iter, should_checkpoint, should_exit, exit_code, grad_norm, num_zeros_in_grad
+    return {}, skipped_iter, should_checkpoint, should_exit, exit_code, grad_norm, num_zeros_in_grad
+
+
+def training_log(loss_dict, total_loss_dict, learning_rate, decoupled_learning_rate, iteration,
+                 loss_scale, report_memory_flag, skipped_iter,
+                 grad_norm, params_norm, num_zeros_in_grad):
+    """Log training information such as losses, timing, ...."""
+    args = get_args()
+    timers = get_timers()
+    writer = get_tensorboard_writer()
+    wandb_writer = get_wandb_writer()
+    one_logger = get_one_logger()
+
+    # Advanced, skipped, and Nan iterations.
+    advanced_iters_key = 'advanced iterations'
+    skipped_iters_key = 'skipped iterations'
+    nan_iters_key = 'nan iterations'
+    # Advanced iterations.
+    if not skipped_iter:
+        total_loss_dict[advanced_iters_key] = total_loss_dict.get(
+            advanced_iters_key, 0) + 1
+    else:
+        if advanced_iters_key not in total_loss_dict:
+            total_loss_dict[advanced_iters_key] = 0
+    # Skipped iterations.
+    total_loss_dict[skipped_iters_key] = total_loss_dict.get(
+        skipped_iters_key, 0) + skipped_iter
+    # Update losses and set nan iterations
+    got_nan = False
+    for key in loss_dict:
+        if not skipped_iter:
+            total_loss_dict[key] = total_loss_dict.get(
+                key, torch.tensor([0.0], dtype=torch.float, device='cuda')) + loss_dict[key]
+        else:
+            value = loss_dict[key].float().sum().item()
+            is_nan = value == float('inf') or \
+                     value == -float('inf') or \
+                     value != value
+            got_nan = got_nan or is_nan
+    total_loss_dict[nan_iters_key] = total_loss_dict.get(
+        nan_iters_key, 0) + int(got_nan)
+
+    # Logging.
+    timers_to_log = [
+        'forward-backward',
+        'forward-compute',
+        'backward-compute',
+        'batch-generator',
+        'forward-recv',
+        'forward-send',
+        'backward-recv',
+        'backward-send',
+        'forward-send-forward-recv',
+        'forward-send-backward-recv',
+        'backward-send-forward-recv',
+        'backward-send-backward-recv',
+        'forward-backward-send-forward-backward-recv',
+        'layernorm-grads-all-reduce',
+        'embedding-grads-all-reduce',
+        'all-grads-sync',
+        'params-all-gather',
+        'optimizer-copy-to-main-grad',
+        'optimizer-unscale-and-check-inf',
+        'optimizer-clip-main-grad',
+        'optimizer-count-zeros',
+        'optimizer-inner-step',
+        'optimizer-copy-main-to-model-params',
+        'optimizer']
+
+    # Calculate batch size.
+    batch_size = args.micro_batch_size * args.data_parallel_size * \
+        get_num_microbatches()
+
+    # Track app tag & app tag ID
+    one_logger_utils.track_app_tag(batch_size, args.world_size, args.seq_length)
+
+    total_iterations = total_loss_dict[advanced_iters_key] + \
+                       total_loss_dict[skipped_iters_key]
+
+    # Tensorboard values.
+    # Timer requires all the ranks to call.
+    if args.log_timers_to_tensorboard and \
+       (iteration % args.tensorboard_log_interval == 0):
+        timers.write(timers_to_log, writer, iteration,
+                     normalizer=total_iterations)
+    if writer and (iteration % args.tensorboard_log_interval == 0):
+        if wandb_writer:
+            wandb_writer.log({'samples vs steps': args.consumed_train_samples},
+                             iteration)
+        writer.add_scalar('learning-rate', learning_rate, iteration)
+        if args.decoupled_lr is not None:
+            writer.add_scalar('decoupled-learning-rate', decoupled_learning_rate, iteration)
+        writer.add_scalar('learning-rate vs samples', learning_rate,
+                          args.consumed_train_samples)
+        if wandb_writer:
+            wandb_writer.log({'learning-rate': learning_rate}, iteration)
+        if args.skipped_train_samples > 0:
+            writer.add_scalar('skipped-train-samples', args.skipped_train_samples, iteration)
+            if wandb_writer:
+                wandb_writer.log({'skipped-train-samples': args.skipped_train_samples}, iteration)
+        writer.add_scalar('batch-size', batch_size, iteration)
+        writer.add_scalar('batch-size vs samples', batch_size,
+                          args.consumed_train_samples)
+        if wandb_writer:
+            wandb_writer.log({'batch-size': batch_size}, iteration)
+        for key in loss_dict:
+            writer.add_scalar(key , loss_dict[key], iteration)
+            writer.add_scalar(key + ' vs samples', loss_dict[key],
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({key: loss_dict[key]}, iteration)
+        if args.log_loss_scale_to_tensorboard:
+            writer.add_scalar('loss-scale', loss_scale, iteration)
+            writer.add_scalar('loss-scale vs samples', loss_scale,
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({'loss-scale': loss_scale}, iteration)
+        if args.log_world_size_to_tensorboard:
+            writer.add_scalar('world-size', args.world_size, iteration)
+            writer.add_scalar('world-size vs samples', args.world_size,
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({'world-size': args.world_size}, iteration)
+        if grad_norm is not None:
+            writer.add_scalar('grad-norm', grad_norm, iteration)
+            writer.add_scalar('grad-norm vs samples', grad_norm,
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({'grad-norm': grad_norm}, iteration)
+        if num_zeros_in_grad is not None:
+            writer.add_scalar('num-zeros', num_zeros_in_grad, iteration)
+            writer.add_scalar('num-zeros vs samples', num_zeros_in_grad,
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({'num-zeros': num_zeros_in_grad}, iteration)
+        if params_norm is not None:
+            writer.add_scalar('params-norm', params_norm, iteration)
+            writer.add_scalar('params-norm vs samples', params_norm,
+                              args.consumed_train_samples)
+            if wandb_writer:
+                wandb_writer.log({'params-norm': params_norm}, iteration)
+        if args.log_memory_to_tensorboard:
+            mem_stats = torch.cuda.memory_stats()
+            writer.add_scalar(
+                "mem-reserved-bytes",
+                mem_stats["reserved_bytes.all.current"],
+                iteration,
+            )
+            writer.add_scalar(
+                "mem-allocated-bytes",
+                mem_stats["allocated_bytes.all.current"],
+                iteration,
+            )
+            writer.add_scalar(
+                "mem-allocated-count",
+                mem_stats["allocation.all.current"],
+                iteration,
+            )
+    if args.num_experts is not None:
+        moe_loss_scale = 1 / get_num_microbatches()
+        track_moe_metrics(moe_loss_scale, iteration, writer, wandb_writer, total_loss_dict, args.moe_per_layer_logging)
+
+    if iteration % args.log_interval == 0:
+        # HACK(huang.huang): support memory analysis dump
+        if args.record_memory_history:
+            snapshot = torch.cuda.memory._snapshot()
+            from pickle import dump
+            os.makedirs("./memory_snapshot", exist_ok=True)
+            with open(f"./memory_snapshot/iter{iteration}-{args.memory_snapshot_path}", 'wb') as f:
+                dump(snapshot, f)
+        ## HACK(huang.huang)
+
+        elapsed_time = timers('interval-time').elapsed(barrier=True)
+        elapsed_time_per_iteration = elapsed_time / total_iterations
+
+        throughput = num_floating_point_operations(args, batch_size) / (
+            elapsed_time_per_iteration * 10**12 * args.world_size)
+
+        one_logger_utils.track_e2e_metrics(args.log_throughput, throughput)
+
+        if args.log_timers_to_tensorboard:
+            if writer:
+                writer.add_scalar('iteration-time',
+                                  elapsed_time_per_iteration, iteration)
+            if wandb_writer:
+                wandb_writer.log({'iteration-time': elapsed_time_per_iteration},
+                                 iteration)
+        log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
+        log_string += ' iteration {:8d}/{:8d} |'.format(
+            iteration, args.train_iters)
+        # mfu = throughput / 465
+        # tokens_per_gpu_per_second = float(batch_size * args.seq_length) / elapsed_time_per_iteration / args.world_size
+        # log_string += ' tokens_per_gpu_per_second: {:.2f} /s |'.format(tokens_per_gpu_per_second)
+        # log_string += ' mfu: {:.4f} |'.format(mfu)
+        log_string += ' consumed samples: {:12d} |'.format(
+            args.consumed_train_samples)
+        if args.skipped_train_samples > 0:
+            log_string += ' skipped samples: {:12d} |'.format(
+                args.skipped_train_samples)
+        log_string += ' elapsed time per iteration (ms): {:.1f} |'.format(
+            elapsed_time_per_iteration * 1000.0)
+        if args.log_throughput:
+            log_string += f' throughput per GPU (TFLOP/s/GPU): {throughput:.1f} |'
+            if args.log_timers_to_tensorboard:
+                if writer:
+                    writer.add_scalar('throughput', throughput, iteration)
+                if wandb_writer:
+                    wandb_writer.log({'throughput': throughput}, iteration)
+        assert learning_rate is not None
+        # Decoupled_learning_rate should be not None only on first and last pipeline stage.
+        log_string += ' learning rate: {:.6E} |'.format(learning_rate)
+        if args.decoupled_lr is not None and (mpu.is_pipeline_first_stage(ignore_virtual=True) or
+                                              mpu.is_pipeline_last_stage(ignore_virtual=True)):
+            assert decoupled_learning_rate is not None
+            log_string += ' decoupled learning rate: {:.6E} |'.format(decoupled_learning_rate)
+        else:
+            assert decoupled_learning_rate is None
+        log_string += ' global batch size: {:5d} |'.format(batch_size)
+        current_loss_dic = dict()
+        for key in total_loss_dict:
+            if key not in [advanced_iters_key, skipped_iters_key,
+                           nan_iters_key]:
+                avg = total_loss_dict[key].item() / \
+                      float(max(1, total_loss_dict[advanced_iters_key]))
+                if avg > 0.0:
+                    log_string += ' {}: {:.6E} |'.format(key, avg)
+                    current_loss_dic[key] = avg
+                total_loss_dict[key] = torch.tensor([0.0], dtype=torch.float, device='cuda')
+        log_string += ' loss scale: {:.1f} |'.format(loss_scale)
+        if grad_norm is not None:
+            log_string += ' grad norm: {:.3f} |'.format(grad_norm)
+        if num_zeros_in_grad is not None:
+            log_string += ' num zeros: {:.1f} |'.format(num_zeros_in_grad)
+        if params_norm is not None:
+            log_string += ' params norm: {:.3f} |'.format(params_norm)
+        log_string += ' number of skipped iterations: {:3d} |'.format(
+            total_loss_dict[skipped_iters_key])
+        log_string += ' number of nan iterations: {:3d} |'.format(
+            total_loss_dict[nan_iters_key])
+        total_loss_dict[advanced_iters_key] = 0
+        total_loss_dict[skipped_iters_key] = 0
+        total_loss_dict[nan_iters_key] = 0
+        print_rank_last(log_string)
+        if report_memory_flag and learning_rate > 0.:
+            # Report memory after optimizer state has been initialized.
+            if torch.distributed.get_rank() == 0:
+                num_microbatches = get_num_microbatches()
+                report_theoretical_memory(args, num_microbatches=num_microbatches, verbose=True)
+            report_memory('(after {} iterations)'.format(iteration))
+            # report_memory_flag = False
+        timers.log(timers_to_log, normalizer=args.log_interval)
+
+        # log to mlflow
+        if need_mlflow():
+            mlflow_metrics = current_loss_dic
+            mlflow_metrics['mfu'] = mfu
+            mlflow_metrics["tps"] = tokens_per_gpu_per_second
+            mlflow_metrics['learning-rate'] = learning_rate
+            mlflow_metrics['consumed-samples'] = args.consumed_train_samples
+            mlflow_metrics['batch-size'] = batch_size
+            mlflow_metrics['loss-scale'] = loss_scale
+            mlflow_metrics['iteration-time'] = elapsed_time_per_iteration
+            mlflow_metrics['world-size'] = args.world_size
+
+            mlflow.log_metrics(mlflow_metrics, step=iteration, synchronous=False)
+
+    return report_memory_flag
+
+def train(forward_step_func, model, optimizer, opt_param_scheduler,
+          train_data_iterator, valid_data_iterator,
+          process_non_loss_data_func, config, checkpointing_context, non_loss_data_func):
+    """Train the model function."""
+    args = get_args()
+    timers = get_timers()
+    one_logger = get_one_logger()
+
+    # Write args to tensorboard
+    write_args_to_tensorboard()
+
+    # Turn on training mode which enables dropout.
+    for model_module in model:
+        model_module.train()
+
+    # Tracking loss.
+    total_loss_dict = {}
+
+    # Iterations.
+    iteration = args.iteration
+
+    # Track E2E metrics at the start of training
+    one_logger_utils.on_train_start(iteration=iteration, consumed_train_samples=args.consumed_train_samples,
+                                    train_samples=args.train_samples, seq_length=args.seq_length,
+                                    train_iters=args.train_iters, save=args.save, async_save=args.async_save,
+                                    log_throughput=args.log_throughput,
+                                    num_floating_point_operations_so_far=args.num_floating_point_operations_so_far)
+
+    num_floating_point_operations_so_far = args.num_floating_point_operations_so_far
+
+    # Setup some training config params
+    config.grad_scale_func = optimizer.scale_loss
+    config.timers = timers
+    if isinstance(model[0], DDP) and args.overlap_grad_reduce:
+        assert config.no_sync_func is None, \
+            ('When overlap_grad_reduce is True, config.no_sync_func must be None; '
+             'a custom no_sync_func is not supported when overlapping grad-reduce')
+        config.no_sync_func = [model_chunk.no_sync for model_chunk in model]
+        if len(model) == 1:
+            config.no_sync_func = config.no_sync_func[0]
+        if args.align_grad_reduce:
+            config.grad_sync_func = [model_chunk.start_grad_sync for model_chunk in model]
+            if len(model) == 1:
+                config.grad_sync_func = config.grad_sync_func[0]
+    if args.overlap_param_gather and args.align_param_gather:
+        config.param_sync_func = [model_chunk.start_param_sync for model_chunk in model]
+        if len(model) == 1:
+            config.param_sync_func = config.param_sync_func[0]
+    config.finalize_model_grads_func = finalize_model_grads
+
+    timers('interval-time', log_level=0).start(barrier=True)
+    print_datetime('before the start of training step')
+    report_memory_flag = True
+    # exit = False
+    pre_hook_enabled = False
+    should_exit = False
+    exit_code = 0
+
+    if args.manual_gc:
+        # Disable the default garbage collector and perform the collection manually.
+        # This is to align the timing of garbage collection across ranks.
+        assert args.manual_gc_interval >= 0, \
+            'Manual garbage collection interval should be laerger than or equal to 0.'
+        gc.disable()
+        gc.collect()
+
+    # Singleton Initialization
+    if args.log_straggler:
+        global stimer
+        world = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+        mmcnt = args.straggler_minmax_count
+        stimer.configure(world, rank,
+                mmcnt = mmcnt,
+                enabled = not args.disable_straggler_on_startup,
+                port = args.straggler_ctrlr_port)
+    # total_flops = 0.0
+    num_floating_point_operations_since_last_log_event = 0.0
+
+    num_microbatches = get_num_microbatches()
+    eval_duration = 0.0
+    eval_iterations = 0
+
+    def get_e2e_base_metrics():
+        """Get base metrics values for one-logger to calculate E2E tracking metrics.
+        """
+        return {
+            'iteration': iteration,
+            'train_duration': timers('interval-time').active_time(),
+            'eval_duration': eval_duration,
+            'eval_iterations': eval_iterations,
+            'total_flops': num_floating_point_operations_since_last_log_event,
+            'num_floating_point_operations_so_far': num_floating_point_operations_so_far,
+            'consumed_train_samples': args.consumed_train_samples,
+            'world_size': args.world_size,
+            'seq_length': args.seq_length
+        }
+    # Cache into one-logger for callback
+    if one_logger:
+        with one_logger.get_context_manager():
+            one_logger.store_set('get_e2e_base_metrics', get_e2e_base_metrics)
+
+    prof = None
+    if args.profile and torch.distributed.get_rank() in args.profile_ranks and args.use_pytorch_profiler:
+        prof = torch.profiler.profile(
+        schedule=torch.profiler.schedule(
+            wait=max(args.profile_step_start-1, 0),
+            warmup=1 if args.profile_step_start > 0 else 0,
+            active=args.profile_step_end-args.profile_step_start,
+            repeat=1),
+        on_trace_ready=torch.profiler.tensorboard_trace_handler(args.tensorboard_dir),
+        record_shapes=True,
+        with_stack=True)
+        prof.start()
+
+    start_iteration = iteration
+    # Disable forward pre-hook to start training to ensure that errors in checkpoint loading
+    # or random initialization don't propagate to all ranks in first all-gather (which is a
+    # no-op if things work correctly).
+    if args.use_distributed_optimizer and args.overlap_param_gather:
+        disable_forward_pre_hook(model, param_sync=False)
+        # Also remove param_sync_func temporarily so that sync calls made in
+        # `forward_backward_func` are no-ops.
+        param_sync_func = config.param_sync_func
+        config.param_sync_func = None
+        pre_hook_enabled = False
+    # Also, check weight hash across DP replicas to be very pedantic.
+    if args.check_weight_hash_across_dp_replicas_interval is not None:
+        assert check_param_hashes_across_dp_replicas(model, cross_check=True), \
+            "Parameter hashes not matching across DP replicas"
+        torch.distributed.barrier()
+        print_rank_0(f">>> Weight hashes match after {iteration} iterations...")
+
+    # HACK(dongsheng.zhang) modelstudio init
+    if need_mlflow():
+        mlflow.start_run()
+
+    with maybe_enable_profiling(
+        args, global_step=iteration
+    ) as torch_profiler:
+        while iteration < args.train_iters:
+            if args.profile and torch.distributed.get_rank() in args.profile_ranks:
+                if args.use_pytorch_profiler:
+                    prof.step()
+                elif iteration == args.profile_step_start:
+                    torch.cuda.cudart().cudaProfilerStart()
+                    torch.autograd.profiler.emit_nvtx(record_shapes=True).__enter__()
+
+            maybe_finalize_async_save(blocking=False)
+
+            # Update number of microbatches first without consistency check to decide if a
+            # checkpoint should be saved. If the number of microbatches is different
+            # from the previous iteration, save a checkpoint. Then run consistency check
+            # to make sure training configuration is still valid.
+            update_num_microbatches(args.consumed_train_samples, consistency_check=False, verbose=True)
+            if get_num_microbatches() != num_microbatches and iteration != 0:
+                assert get_num_microbatches() > num_microbatches, \
+                    "number of microbatches should be increasing due to batch size rampup ... %d -> %d." % (num_microbatches, get_num_microbatches())
+                if args.save is not None:
+                    save_checkpoint_and_time(iteration, model, optimizer,
+                                            opt_param_scheduler,
+                                            num_floating_point_operations_so_far,
+                                            checkpointing_context, train_data_iterator=train_data_iterator)
+            num_microbatches = get_num_microbatches()
+            update_num_microbatches(args.consumed_train_samples, consistency_check=True, verbose=True)
+
+            args.curr_iteration = iteration
+            loss_dict, skipped_iter, should_checkpoint, should_exit, exit_code, grad_norm, num_zeros_in_grad = \
+                train_step(forward_step_func,
+                        train_data_iterator,
+                        model,
+                        optimizer,
+                        opt_param_scheduler,
+                        config)
+
+            if should_checkpoint:
+                save_checkpoint_and_time(iteration, model, optimizer,
+                                        opt_param_scheduler,
+                                        num_floating_point_operations_so_far,
+                                        checkpointing_context, train_data_iterator=train_data_iterator)
+            if should_exit:
+                break
+
+            # Enable forward pre-hooks after first set of forward and backward passes.
+            # When running in fp16, skip all NaN iterations until steady-state loss scaling value
+            # is reached.
+            if iteration == start_iteration:
+                if skipped_iter:
+                    # Only enable forward pre-hook after a training step has successfully run. Relevant
+                    # for fp16 codepath where first XX iterations are skipped until steady-state loss
+                    # scale value is reached.
+                    start_iteration = iteration + 1
+                else:
+                    # Enable forward pre-hook after training step has successfully run. All subsequent
+                    # forward passes will use the forward pre-hook / `param_sync_func` in
+                    # `forward_backward_func`.
+                    if args.use_distributed_optimizer and args.overlap_param_gather:
+                        enable_forward_pre_hook(model)
+                        config.param_sync_func = param_sync_func
+                        pre_hook_enabled = True
+
+            if torch_profiler:
+                torch_profiler.step()
+            iteration += 1
+            batch_size = mpu.get_data_parallel_world_size() * \
+                        args.micro_batch_size * \
+                        get_num_microbatches()
+            args.consumed_train_samples += batch_size
+            num_skipped_samples_in_batch = (get_current_global_batch_size() -
+                                            get_current_running_global_batch_size())
+            if args.decrease_batch_size_if_needed:
+                assert num_skipped_samples_in_batch >= 0
+            else:
+                assert num_skipped_samples_in_batch == 0
+            args.skipped_train_samples += num_skipped_samples_in_batch
+            num_floating_point_operations_in_batch = num_floating_point_operations(args, batch_size)
+            num_floating_point_operations_so_far += num_floating_point_operations_in_batch
+            num_floating_point_operations_since_last_log_event += num_floating_point_operations_in_batch
+
+            # # Send heartbeat to FT package and update timeouts.
+            # if args.enable_ft_package:
+            #     ft_client = ft_integration.get_rank_monitor_client(
+            #         ft_integration.StateMachineActions.TRAIN_HEARTBEAT)
+            #     if ft_client is not None:
+            #         ft_client.send_heartbeat()
+            #         # TODO we are always calculating timeouts in the current implementation
+            #         # if we want to rely on manually setup then we need to add additional argument
+            #         # to training and pass it here
+            #         if ft_integration.can_update_timeouts():
+            #             ft_integration.get_rank_monitor_client(
+            #                 ft_integration.StateMachineActions.UPDATE_TIMEOUT).calculate_and_set_timeouts()
+            #             print_rank_0(f'Updated FT timeouts. New values: \
+            #                 {ft_integration.get_rank_monitor_client().timeouts}')
+
+            # # Bring CPU and GPU back in sync if on right iteration.
+            # if (
+            #     args.train_sync_interval
+            #     and iteration % args.train_sync_interval == 0
+            # ):
+            #     torch.cuda.synchronize()
+
+            # Logging.
+            if not optimizer.is_stub_optimizer:
+                loss_scale = optimizer.get_loss_scale().item()
+            else:
+                loss_scale = 1.0
+            params_norm = None
+            if args.log_params_norm:
+                params_norm = calc_params_l2_norm(model)
+
+            learning_rate = None
+            decoupled_learning_rate = None
+            for param_group in optimizer.param_groups:
+                if param_group['is_decoupled_lr']:
+                    decoupled_learning_rate = param_group['lr']
+                else:
+                    learning_rate = param_group['lr']
+            report_memory_flag = training_log(loss_dict, total_loss_dict,
+                                            learning_rate,
+                                            decoupled_learning_rate,
+                                            iteration, loss_scale,
+                                            report_memory_flag, skipped_iter,
+                                            grad_norm, params_norm, num_zeros_in_grad)
+
+            # # StragglerDetector
+            # if iteration % args.log_interval == 0 and args.log_straggler:
+            #     stimer.report(total_flops, args.log_interval)
+            #     total_flops = 0.0
+
+            # if args.check_weight_hash_across_dp_replicas_interval is not None and \
+            #         iteration % args.check_weight_hash_across_dp_replicas_interval == 0:
+            #     if args.use_distributed_optimizer and args.overlap_param_gather:
+            #         optimizer.disable_pre_hook()
+            #     assert check_param_hashes_across_dp_replicas(model, cross_check=True), \
+            #         "Parameter hashes not matching across DP replicas"
+            #     torch.distributed.barrier()
+            #     print_rank_0(f">>> Weight hashes match after {iteration} iterations...")
+            #     if args.use_distributed_optimizer and args.overlap_param_gather:
+            #         optimizer.enable_pre_hook()
+
+            # # Autoresume
+            # if args.adlr_autoresume and \
+            # (iteration % args.adlr_autoresume_interval == 0):
+            #     check_adlr_autoresume_termination(iteration, model, optimizer,
+            #                                     opt_param_scheduler)
+
+            # Evaluation
+            if args.eval_interval and iteration % args.eval_interval == 0 and \
+                args.do_valid:
+                timers('interval-time').stop()
+                if args.use_distributed_optimizer and args.overlap_param_gather:
+                    disable_forward_pre_hook(model)
+                    pre_hook_enabled = False
+                if args.manual_gc and args.manual_gc_eval:
+                    # Collect all objects.
+                    gc.collect()
+                prefix = f'iteration {iteration}'
+                timers('eval-time', log_level=0).start(barrier=True)
+                evaluate_and_print_results(prefix, forward_step_func,
+                                        valid_data_iterator, model,
+                                        iteration, process_non_loss_data_func,
+                                        config, verbose=False, write_to_tensorboard=True,
+                                        non_loss_data_func=non_loss_data_func)
+                eval_duration += timers('eval-time').elapsed()
+                eval_iterations += args.eval_iters
+                timers('eval-time').stop()
+                one_logger_utils.track_e2e_metrics()
+
+                if args.manual_gc and args.manual_gc_eval:
+                    # Collect only the objects created and used in evaluation.
+                    gc.collect(generation=0)
+                if args.use_distributed_optimizer and args.overlap_param_gather:
+                    enable_forward_pre_hook(model)
+                    pre_hook_enabled = True
+                timers('interval-time', log_level=0).start(barrier=True)
+
+
+                if args.enable_ft_package and ft_integration.get_rank_monitor_client() is not None:
+                    ft_integration.get_rank_monitor_client(
+                        ft_integration.StateMachineActions.EVAL_HEARTBEAT).send_heartbeat()
+
+            # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC).
+            # Some of these only happen at specific iterations.
+            post_training_step_callbacks(model, optimizer, opt_param_scheduler, iteration, prof,
+                                        num_floating_point_operations_since_last_log_event)
+
+            # Checkpoint and decide whether to exit.
+            should_exit = checkpoint_and_decide_exit(model, optimizer, opt_param_scheduler, iteration,
+                                                    num_floating_point_operations_so_far,
+                                                    checkpointing_context, train_data_iterator)
+            if should_exit:
+                break
+
+    one_logger_utils.track_e2e_metrics()
+
+    # Flush TensorBoard, WandB writers and one-logger.
+    writer = get_tensorboard_writer()
+    if writer:
+        writer.flush()
+
+    # Close out pre-hooks if using distributed optimizer and overlapped param gather.
+    if pre_hook_enabled:
+        disable_forward_pre_hook(model)
+
+    if args.enable_ft_package and ft_integration.get_rank_monitor_client() is not None:
+        ft_integration.get_rank_monitor_client().shutdown_workload_monitoring()
+
+    maybe_finalize_async_save(blocking=True)
+
+    # If any exit conditions (signal handler, duration, iterations) have been reached, exit.
+    if should_exit:
+        wandb_writer = get_wandb_writer()
+        if wandb_writer:
+            wandb_writer.finish()
+        sys.exit(exit_code)
+
+    return iteration, num_floating_point_operations_so_far
+
+
+import megatron.training
+megatron.training.training.train_step = train_step
+megatron.training.training.training_log = training_log
+megatron.training.training.train = train
diff --git a/megatron-lm-musa-patch/musa_patch/utils.py b/megatron-lm-musa-patch/musa_patch/utils.py
new file mode 100644
index 000000000..ab83ec741
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/utils.py
@@ -0,0 +1,68 @@
+import contextlib
+import gc
+import math
+import os
+import subprocess
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Generator, Iterable, List, Optional, Set, Union
+
+# hardcoded BF16 type peak flops for Moore Threads S4000, and S5000 GPU
+def get_peak_flops(device_name: str) -> int:
+    try:
+        # Run the lspci command and capture the output
+        result = subprocess.run(["lspci", "-d", "1ed5:"], stdout=subprocess.PIPE, text=True)
+        # Filter the output for lines containing both "NVIDIA" and "H100"
+        filtered_lines = [
+            line
+            for line in result.stdout.splitlines()
+            if "NVIDIA" in line and "H100" in line
+        ]
+        # Join all filtered lines into a single string
+        device_name = " ".join(filtered_lines) or device_name
+    except FileNotFoundError as e:
+        logger.warning(f"Error running lspci: {e}, fallback to use device_name")
+    if "A100" in device_name:
+        # data from https://www.nvidia.com/en-us/data-center/a100/
+        return 312e12
+    elif "H100" in device_name:
+        # data from https://www.nvidia.com/en-us/data-center/h100/
+        # NOTE: Specifications are one-half lower without sparsity.
+        if "NVL" in device_name:
+            return 835e12
+        elif "PCIe" in device_name:
+            return 756e12
+        else:  # for H100 SXM and other variants
+            return 989e12
+    elif "H200" in device_name:
+        # data from https://www.nvidia.com/en-us/data-center/h200/
+        return 989e12
+    else:  # for other GPU types, assume A100
+        logger.warning(f"Peak flops undefined for: {device_name}, fallback to A100")
+        return 312e12
+
+
+@dataclass(frozen=True)
+class Color:
+    black = "\033[30m"
+    red = "\033[31m"
+    green = "\033[32m"
+    yellow = "\033[33m"
+    blue = "\033[34m"
+    magenta = "\033[35m"
+    cyan = "\033[36m"
+    white = "\033[37m"
+    reset = "\033[39m"
+
+
+@dataclass(frozen=True)
+class NoColor:
+    black = ""
+    red = ""
+    green = ""
+    yellow = ""
+    blue = ""
+    magenta = ""
+    cyan = ""
+    white = ""
+    reset = ""
diff --git a/megatron-lm-musa-patch/musa_patch/zbb_light/__init__.py b/megatron-lm-musa-patch/musa_patch/zbb_light/__init__.py
new file mode 100644
index 000000000..8b0d05a82
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/zbb_light/__init__.py
@@ -0,0 +1,2 @@
+
+from .zbpp_light import patch_megatron
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/zbb_light/auto_schedule.py b/megatron-lm-musa-patch/musa_patch/zbb_light/auto_schedule.py
new file mode 100644
index 000000000..9772df82d
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/zbb_light/auto_schedule.py
@@ -0,0 +1,93 @@
+from dataclasses import dataclass
+
+@dataclass(eq=True)
+class ScheduledNode:
+    type: str
+    stage: int
+    minibatch: int
+
+def auto_schedule(nstages, nmb):
+    f = [0] * nstages
+    b = [0] * nstages
+    # W is a stack for each stage
+    w = [[] for i in range(nstages)]
+    result = [[] for i in range(nstages)]
+    last_compute_id = [-1] * nstages
+    def schedule_f(stage):
+        if not stage == 0:
+            result[stage].append(ScheduledNode(
+                'RECV_FORWARD',
+                stage=stage,
+                minibatch=f[stage]))
+        result[stage].append(
+            ScheduledNode(
+              type='F',
+              stage=stage,
+              minibatch=f[stage]))
+        last_compute_id[stage] = len(result[stage]) - 1
+        if not stage == nstages - 1:
+            result[stage].append(ScheduledNode(
+                'SEND_FORWARD',
+                stage=stage,
+                minibatch=f[stage]))
+        f[stage] += 1
+    def schedule_b(stage):
+        if not stage == nstages - 1:
+            result[stage].append(ScheduledNode(
+                'RECV_BACKWARD',
+                stage=stage,
+                minibatch=b[stage]))
+        result[stage].append(
+            ScheduledNode(
+              type='B',
+              stage=stage,
+              minibatch=b[stage]))
+        last_compute_id[stage] = len(result[stage]) - 1
+        if not stage == 0:
+            result[stage].append(ScheduledNode(
+                'SEND_BACKWARD',
+                stage=stage,
+                minibatch=b[stage]))
+        w[stage].append(b[stage])
+        b[stage] += 1
+    def schedule_w(stage):
+        assert last_compute_id[stage] != -1
+        if result[stage][last_compute_id[stage]].type == 'B' and (stage == 0 or not f[stage] == nmb):
+            result[stage][last_compute_id[stage]].type = 'BW'
+        else:
+            result[stage].append(
+                ScheduledNode(
+                  type='W',
+                  stage=stage,
+                  minibatch=w[stage][-1]))
+            last_compute_id[stage] = len(result[stage]) - 1
+        w[stage].pop(-1)
+
+
+    for stage in range(nstages):
+        num_warmup_microbatches = nstages - stage - 1
+        num_warmup_microbatches = min(num_warmup_microbatches, nmb)
+        remaining = nmb - num_warmup_microbatches
+        for i in range(num_warmup_microbatches):
+            schedule_f(stage)
+        for i in range(remaining):
+            schedule_f(stage)
+            schedule_b(stage)
+            if i >= stage:
+                schedule_w(stage)
+        for i in range(num_warmup_microbatches):
+            schedule_b(stage)
+            if remaining + i >= stage:
+                schedule_w(stage)
+        assert f[stage] == b[stage] == nmb
+        while len(w[stage]) > 0:
+            schedule_w(stage)
+
+    for stage in range(nstages):
+        print(' '.join([f'{x.type}{x.minibatch}' for x in result[stage] if x.type in {'F', 'B', 'W', 'BW'}]))
+    return result
+
+
+if __name__ == "__main__":
+    auto_schedule(4, 12)
+    auto_schedule(8, 4)
diff --git a/megatron-lm-musa-patch/musa_patch/zbb_light/weight_grad_store.py b/megatron-lm-musa-patch/musa_patch/zbb_light/weight_grad_store.py
new file mode 100644
index 000000000..530b17d1c
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/zbb_light/weight_grad_store.py
@@ -0,0 +1,56 @@
+import queue
+import torch
+from typing import List
+
+class WeightGradStore:
+
+    cache = []
+    weight_grad_queue = queue.Queue()
+    combine_bw = False
+
+    @classmethod
+    def set_combine_bw(cls, combine_bw):
+        # For the following backward pass, combine W with B and skip next W
+        cls.combine_bw = combine_bw
+
+    @classmethod
+    def split_bw(cls):
+        # For the following backward pass, combine W with B and skip next W
+        return not cls.combine_bw
+
+    @classmethod
+    def put(cls, inputs, func, pos_func = None):
+        if cls.combine_bw:
+            func(*inputs)
+            if pos_func is not None:
+                if isinstance(inputs[0], List):
+                    pos_func(*inputs[0])
+                else:
+                    pos_func(inputs[0])
+            return
+        # Store the weight gradient computation of linear layers.
+        cls.cache.append((inputs, func, pos_func))
+
+    @classmethod
+    def flush(cls):
+        if cls.combine_bw:
+            cls.combine_bw = False
+            return
+        # Collect all stored computations during backward as a W.
+        cls.weight_grad_queue.put(cls.cache)
+        cls.cache = []
+
+    @classmethod
+    def pop(cls):
+        assert not cls.combine_bw
+        # Execute a single W.
+        assert cls.weight_grad_queue.qsize() > 0
+        stored_grads = cls.weight_grad_queue.get()
+        with torch.enable_grad():
+            for inputs, func, pos_func in stored_grads:
+                func(*inputs)
+                if pos_func is not None:
+                    if isinstance(inputs[0], List):
+                        pos_func(*inputs[0])
+                    else:
+                        pos_func(inputs[0])
\ No newline at end of file
diff --git a/megatron-lm-musa-patch/musa_patch/zbb_light/zb_schedule.py b/megatron-lm-musa-patch/musa_patch/zbb_light/zb_schedule.py
new file mode 100644
index 000000000..847f3d54c
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/zbb_light/zb_schedule.py
@@ -0,0 +1,580 @@
+import contextlib
+import itertools
+from typing import Iterator, List, Union
+import os
+import torch
+
+from megatron import core
+from megatron.core import parallel_state
+from megatron.core.utils import get_model_config, get_model_type, get_model_xattn
+from megatron.core.parallel_state import (
+    get_pipeline_model_parallel_group,
+    get_pipeline_model_parallel_next_rank,
+    get_pipeline_model_parallel_prev_rank,
+)
+from megatron.core.pipeline_parallel.schedules import (
+    recv_forward,
+    send_forward,
+    recv_backward,
+    send_backward,
+    deallocate_output_tensor,
+    forward_step,
+    backward_step,
+    get_tensor_shapes,
+)
+from megatron.core.num_microbatches_calculator import get_num_microbatches
+from . import auto_schedule
+from .weight_grad_store import WeightGradStore
+
+
+AUTO_SCHEDULE_COMMUNICATION_TYPES = {'RECV_FORWARD', 'RECV_BACKWARD', 'SEND_FORWARD', 'SEND_BACKWARD'}
+
+
+def fused_pipeline_ops(
+    tensor_send_prev: List[torch.Tensor],
+    tensor_recv_prev: List[torch.Tensor],
+    tensor_send_next: List[torch.Tensor],
+    tensor_recv_next: List[torch.Tensor],
+    group: torch.distributed.ProcessGroup,
+):
+    ops = []
+    if parallel_state.get_pipeline_model_parallel_rank() % 2 == 0:
+        for t in tensor_send_next:
+            send_next_op = torch.distributed.P2POp(
+                torch.distributed.isend,
+                t,
+                get_pipeline_model_parallel_next_rank(),
+                group,
+            )
+            ops.append(send_next_op)
+        for t in tensor_recv_prev:
+            recv_prev_op = torch.distributed.P2POp(
+                torch.distributed.irecv,
+                t,
+                get_pipeline_model_parallel_prev_rank(),
+                group,
+            )
+            ops.append(recv_prev_op)
+        for t in tensor_send_prev:
+            send_prev_op = torch.distributed.P2POp(
+                torch.distributed.isend,
+                t,
+                get_pipeline_model_parallel_prev_rank(),
+                group,
+            )
+            ops.append(send_prev_op)
+        for t in tensor_recv_next:
+            recv_next_op = torch.distributed.P2POp(
+                torch.distributed.irecv,
+                t,
+                get_pipeline_model_parallel_next_rank(),
+                group,
+            )
+            ops.append(recv_next_op)
+    else:
+        for t in tensor_recv_prev:
+            recv_prev_op = torch.distributed.P2POp(
+                torch.distributed.irecv,
+                t,
+                get_pipeline_model_parallel_prev_rank(),
+                group,
+            )
+            ops.append(recv_prev_op)
+        for t in tensor_send_next:
+            send_next_op = torch.distributed.P2POp(
+                torch.distributed.isend,
+                t,
+                get_pipeline_model_parallel_next_rank(),
+                group,
+            )
+            ops.append(send_next_op)
+        for t in tensor_recv_next:
+            recv_next_op = torch.distributed.P2POp(
+                torch.distributed.irecv,
+                t,
+                get_pipeline_model_parallel_next_rank(),
+                group,
+            )
+            ops.append(recv_next_op)
+        for t in tensor_send_prev:
+            send_prev_op = torch.distributed.P2POp(
+                torch.distributed.isend,
+                t,
+                get_pipeline_model_parallel_prev_rank(),
+                group,
+            )
+            ops.append(send_prev_op)
+
+    if len(ops) > 0:
+        reqs = torch.distributed.batch_isend_irecv(ops)
+        for req in reqs:
+            req.wait()
+    else:
+        reqs = []
+    return reqs,[],[],[],[]
+
+def p2p_pipeline_ops(
+    tensor_send_prev: List[torch.Tensor],
+    tensor_recv_prev: List[torch.Tensor],
+    tensor_send_next: List[torch.Tensor],
+    tensor_recv_next: List[torch.Tensor],
+    group: torch.distributed.ProcessGroup,
+):
+    assert( len(tensor_send_prev) <= 1 and len(tensor_recv_prev) <= 1 and len(tensor_send_next)<=1 and len(tensor_recv_next) <=1)
+    reqs, sp_reqs, rp_reqs, sn_reqs, rn_reqs= [], [], [], [], []
+    if parallel_state.get_pipeline_model_parallel_rank() % 2 == 0:
+        for t in tensor_send_next:
+            send_next_req = torch.distributed.isend(
+                tensor=t,
+                dst=get_pipeline_model_parallel_next_rank(),
+                group=group,
+            )
+            sn_reqs.append(send_next_req)
+            reqs.append(send_next_req)
+
+        for t in tensor_recv_prev:
+            recv_prev_req = torch.distributed.irecv(
+                tensor=t,
+                src=get_pipeline_model_parallel_prev_rank(),
+                group=group,
+            )
+            rp_reqs.append(recv_prev_req)
+            reqs.append(recv_prev_req)
+
+        for t in tensor_send_prev:
+            send_prev_req = torch.distributed.isend(
+                tensor=t,
+                dst=get_pipeline_model_parallel_prev_rank(),
+                group=group,
+            )
+            sp_reqs.append(send_prev_req)
+            reqs.append(send_prev_req)
+
+        for t in tensor_recv_next:
+            recv_next_req = torch.distributed.irecv(
+                tensor=t,
+                src=get_pipeline_model_parallel_next_rank(),
+                group=group,
+            )
+            rn_reqs.append(recv_next_req)
+            reqs.append(recv_next_req)
+    else:
+
+        for t in tensor_recv_prev:
+            recv_prev_req = torch.distributed.irecv(
+                tensor=t,
+                src=get_pipeline_model_parallel_prev_rank(),
+                group=group,
+            )
+            rp_reqs.append(recv_prev_req)
+            reqs.append(recv_prev_req)
+
+        for t in tensor_send_next:
+            send_next_req = torch.distributed.isend(
+                tensor=t,
+                dst=get_pipeline_model_parallel_next_rank(),
+                group=group,
+            )
+            sn_reqs.append(send_next_req)
+            reqs.append(send_next_req)
+
+        for t in tensor_recv_next:
+            recv_next_req = torch.distributed.irecv(
+                tensor=t,
+                src=get_pipeline_model_parallel_next_rank(),
+                group=group,
+            )
+            rn_reqs.append(recv_next_req)
+            reqs.append(recv_next_req)
+
+        for t in tensor_send_prev:
+            send_prev_req = torch.distributed.isend(
+                tensor=t,
+                dst=get_pipeline_model_parallel_prev_rank(),
+                group=group,
+            )
+            sp_reqs.append(send_prev_req)
+            reqs.append(send_prev_req)
+
+    # for req in reqs:
+    #         req.wait()
+    return (reqs, sp_reqs, rp_reqs, sn_reqs, rn_reqs)
+
+
+
+def multi_pipeline_ops(
+    tensor_send_prev: List[torch.Tensor],
+    tensor_recv_prev: List[torch.Tensor],
+    tensor_send_next: List[torch.Tensor],
+    tensor_recv_next: List[torch.Tensor],
+):
+    group = get_pipeline_model_parallel_group()
+    if True:
+        p2p_func = fused_pipeline_ops
+    else:
+        p2p_func = p2p_pipeline_ops
+    return p2p_func(
+        tensor_send_prev=tensor_send_prev,
+        tensor_recv_prev=tensor_recv_prev,
+        tensor_send_next=tensor_send_next,
+        tensor_recv_next=tensor_recv_next,
+        group=group,
+    )
+
+
+class ZeroBubbleScheduler:
+
+    def __init__(self):
+        self._reset()
+
+        self.schedules = None
+        self.send_tensor_shapes = None
+        self.recv_tensor_shapes = None
+        self.config = None
+        self.forward_step_func = None
+        self.data_iterator = None
+        self.model = None
+        self.model_type = None
+        self.num_microbatches = None
+        self.collect_non_loss_data = None
+        self.forward_only = None
+        self.no_sync_context = None
+        self.no_sync_func = None
+
+        self.do_post_validation = False
+        self.is_first_run = True
+        self.optimizer = None
+
+    def _free_buffers(self):
+        self.input_tensors = []
+        self.output_tensors = []
+        self.send_forward_buffer = []
+        self.recv_forward_buffer = []
+        self.send_backward_buffer = []
+        self.recv_backward_buffer = []
+        self.forward_data_store = []
+
+    def _reset(self):
+        # Input, output tensors only need to be saved when doing backward passes
+        self._free_buffers()
+        self.send_handles = []
+        self.communication_batch = {
+            'SEND_NEXT': [],
+            'RECV_NEXT': [],
+            'SEND_PREV': [],
+            'RECV_PREV': [],
+        }
+
+    def get_schedules(self):
+        if self.schedules is None:
+            # bootstrap_p2p_communication(self.config)
+            self.schedules = auto_schedule.auto_schedule(
+                parallel_state.get_pipeline_model_parallel_world_size(),
+                get_num_microbatches())[parallel_state.get_pipeline_model_parallel_rank()]
+
+        return self.schedules
+
+    @classmethod
+    def direction_map(cls, node):
+        return {
+            'SEND_FORWARD': 'SEND_NEXT',
+            'RECV_FORWARD': 'RECV_PREV',
+            'SEND_BACKWARD': 'SEND_PREV',
+            'RECV_BACKWARD': 'RECV_NEXT',
+        }[node.type]
+
+    def buffer_map(self, node):
+        return {
+            'SEND_FORWARD': self.send_forward_buffer,
+            'RECV_FORWARD': self.recv_forward_buffer,
+            'SEND_BACKWARD': self.send_backward_buffer,
+            'RECV_BACKWARD': self.recv_backward_buffer,
+        }[node.type]
+
+    def flush(self):
+        name = '_'.join(
+            [f'{v[0].type}.{v[0].minibatch}' for v in itertools.chain(
+                *[vs for k, vs in self.communication_batch.items()])])
+        assert self.send_tensor_shapes == self.recv_tensor_shapes
+        assert len(self.send_tensor_shapes) == 1
+        sn_tensors = [
+            self.buffer_map(x[0]).pop(0)[0]
+            for x in self.communication_batch['SEND_NEXT']
+        ]
+        sp_tensors = [
+            self.buffer_map(x[0]).pop(0)[0]
+            for x in self.communication_batch['SEND_PREV']
+        ]
+
+        rn_tensors = [
+            torch.empty(
+                self.send_tensor_shapes[0],
+                requires_grad=True,
+                device=torch.cuda.current_device(),
+                dtype=self.config.pipeline_dtype,
+            ) for x in self.communication_batch['RECV_NEXT']
+        ]
+        rp_tensors = [
+            torch.empty(
+                self.send_tensor_shapes[0],
+                requires_grad=True,
+                device=torch.cuda.current_device(),
+                dtype=self.config.pipeline_dtype,
+            ) for x in self.communication_batch['RECV_PREV']
+        ]
+        (reqs, sp_reqs, rp_reqs, sn_reqs, rn_reqs) = multi_pipeline_ops(
+            sp_tensors,
+            rp_tensors,
+            sn_tensors,
+            rn_tensors
+        )
+        # We don't care about the reqs order here, all users need to all reqs to finish
+        for x in self.communication_batch['RECV_NEXT']:
+            self.buffer_map(x[0]).append(([rn_tensors.pop(0)], [rn_reqs]))
+        for x in self.communication_batch['RECV_PREV']:
+            self.buffer_map(x[0]).append(([rp_tensors.pop(0)], [rp_reqs]))
+        self.send_handles.append([sp_reqs, sn_reqs])
+        assert(not rn_tensors)
+        assert(not rp_tensors)
+        for direction in ['SEND_PREV', 'SEND_NEXT']:
+            for id, x in enumerate(self.communication_batch[direction]):
+                if x[0].type == 'SEND_FORWARD':
+                    deallocate_output_tensor(sp_tensors[id] if direction == 'SEND_PREV' else sn_tensors[id],
+                                             self.config.deallocate_pipeline_outputs)
+        for k, v in self.communication_batch.items():
+            v.clear()
+
+    def add_communication(
+        self,
+        scheduled_node: auto_schedule.ScheduledNode,
+        next_is_comm: bool,
+        next_compute: auto_schedule.ScheduledNode
+    ):
+        if self.forward_only and 'BACKWARD' in scheduled_node.type:
+            return
+        self.communication_batch[self.direction_map(scheduled_node)].append(
+            (scheduled_node, None))
+        def is_consumer(scheduled_node, next_compute):
+            if scheduled_node.minibatch == next_compute.minibatch:
+                if scheduled_node.type == 'RECV_FORWARD' and next_compute.type == 'F':
+                    return True
+                if scheduled_node.type == 'RECV_BACKWARD' and next_compute.type == 'B':
+                    return True
+            return False
+        if (next_compute is not None and is_consumer(scheduled_node, next_compute)) or not next_is_comm or self.forward_only:
+            self.flush()
+
+    def schedule_f(self, scheduled_node):
+        if core.parallel_state.is_pipeline_first_stage():
+            input_tensor = [None] * len(self.recv_tensor_shapes)
+        else:
+            input_tensor = self.recv_forward_buffer.pop(0)
+            for h in input_tensor[1]:
+                for hh in h:
+                     hh.wait()
+            input_tensor = input_tensor[0]
+
+        output_tensor, _ = forward_step(
+            self.forward_step_func,
+            self.data_iterator,
+            self.model,
+            self.num_microbatches,
+            input_tensor,
+            self.forward_data_store,
+            self.config,
+            self.collect_non_loss_data,
+            checkpoint_activations_microbatch=None,
+        )
+        if not core.parallel_state.is_pipeline_last_stage():
+            self.send_forward_buffer.append(output_tensor)
+        if not self.forward_only:
+            self.input_tensors.append(input_tensor)
+            self.output_tensors.append(output_tensor)
+            if core.parallel_state.is_pipeline_last_stage():
+                deallocate_output_tensor(output_tensor[0], self.config.deallocate_pipeline_outputs)
+
+    def schedule_b(self, scheduled_node):
+        WeightGradStore.set_combine_bw(scheduled_node.type == 'BW')
+        if not self.forward_only:
+            input_tensor = self.input_tensors.pop(0)
+            output_tensor = self.output_tensors.pop(0)
+
+            if core.parallel_state.is_pipeline_last_stage():
+                # Keep the original behavior when we do a dummy communication
+                output_tensor_grad = [None] * len(self.send_tensor_shapes)
+            else:
+                output_tensor_grad = self.recv_backward_buffer.pop(0)
+                for h in output_tensor_grad[1]:
+                    for hh in h:
+                        hh.wait()
+                output_tensor_grad = output_tensor_grad[0]
+            input_tensor_grad = backward_step(
+                input_tensor, output_tensor, output_tensor_grad, self.model_type,
+                self.config
+            )
+            self.send_backward_buffer.append(input_tensor_grad)
+            WeightGradStore.flush()
+
+    def schedule_w(self, scheduled_node):
+        if not self.forward_only:
+            WeightGradStore.pop()
+
+    def disable_grad_sync(self):
+        """Disable asynchronous grad reductions"""
+        if self.no_sync_context is None:
+            self.no_sync_context = self.no_sync_func()
+            self.no_sync_context.__enter__()
+
+    def enable_grad_sync(self):
+        """Enable asynchronous grad reductions"""
+        if self.no_sync_context is not None:
+            self.no_sync_context.__exit__(None, None, None)
+            self.no_sync_context = None
+
+    def prepare(
+        self,
+        forward_step_func,
+        data_iterator: Union[Iterator, List[Iterator]],
+        model: Union[torch.nn.Module, List[torch.nn.Module]],
+        num_microbatches: int,
+        seq_length: int,
+        micro_batch_size: int,
+        decoder_seq_length: int = None,
+        forward_only: bool = False,
+        collect_non_loss_data: bool = False,
+    ):
+        if isinstance(model, list):
+            assert (
+                len(model) == 1
+            ), "non-interleaved pipeline parallelism does not support model chunking"
+            model = model[0]
+        if isinstance(data_iterator, list):
+            assert (
+                len(data_iterator) == 1
+            ), "non-pipeline-parallel schedule does not support model chunking"
+            data_iterator = data_iterator[0]
+
+        config = get_model_config(model)
+        if config.overlap_p2p_comm:
+            raise ValueError(
+                "Non-interleaved pipeline parallelism does not support overlapping p2p communication"
+            )
+        # Disable async grad reductions
+        no_sync_func = config.no_sync_func
+        if no_sync_func is None:
+            no_sync_func = contextlib.nullcontext
+        self.no_sync_func = no_sync_func
+        self.no_sync_context = None
+
+        # Checkpoint the activations of partial Transformer layers in a number of micro-batches
+        # within the maximum outstanding micro-batch backpropagations.
+        # Micro-batches with the ids less than 'num_microbatches_with_partial_activation_checkpoints'
+        # checkpoint partial Transformer layers (or skip checkpointing) and
+        # the rest of micro-batches within a window of micro-batches checkpoint
+        # all Transformer layers. The window of micro-batches is set by the maximum
+        # outstanding backpropagations and becomes smaller at later pipeline stages.
+        # Please refer the appendix C in https://arxiv.org/pdf/2205.05198.pdf
+        assert config.num_microbatches_with_partial_activation_checkpoints is None
+
+        model_type = get_model_type(model)
+        encoder_decoder_xattn = get_model_xattn(model)
+
+        rank = parallel_state.get_pipeline_model_parallel_rank()
+        recv_tensor_shapes = get_tensor_shapes(
+            rank=rank - 1,
+            model_type=model_type,
+            seq_length=seq_length,
+            micro_batch_size=micro_batch_size,
+            decoder_seq_length=decoder_seq_length,
+            config=config,
+            encoder_decoder_xattn=encoder_decoder_xattn,
+
+        )
+        send_tensor_shapes = get_tensor_shapes(
+            rank=rank,
+            model_type=model_type,
+            seq_length=seq_length,
+            micro_batch_size=micro_batch_size,
+            decoder_seq_length=decoder_seq_length,
+            config=config,
+            encoder_decoder_xattn=encoder_decoder_xattn,
+        )
+
+        self.config = config
+        self.model_type = model_type
+        self.recv_tensor_shapes = recv_tensor_shapes
+        self.send_tensor_shapes = send_tensor_shapes
+        self.forward_step_func = forward_step_func
+        self.data_iterator = data_iterator
+        self.model = model
+        self.num_microbatches = num_microbatches
+        self.collect_non_loss_data = collect_non_loss_data
+        self.forward_only = forward_only
+        self._reset()
+        self.it = 0
+
+
+
+    def run(self):
+        # print('-----run:--')
+        schedules = self.get_schedules()
+        self.disable_grad_sync()
+        for it in range(len(schedules)):
+            scheduled_node = schedules[it]
+            # print('----scheduled_node.type:', scheduled_node.type)
+            if scheduled_node.type in AUTO_SCHEDULE_COMMUNICATION_TYPES:
+                next_is_comm = it + 1 < len(schedules) and schedules[it + 1].type in AUTO_SCHEDULE_COMMUNICATION_TYPES
+                next_compute = list(filter(lambda x: x.type in ['F', 'B', 'W'], schedules[it + 1:]))
+                next_compute = next_compute[0] if len(next_compute) > 0 else None
+                self.add_communication(scheduled_node, next_is_comm, next_compute)
+            elif scheduled_node.type == 'F':
+                self.schedule_f(scheduled_node)
+            elif scheduled_node.type in {'B', 'BW'}:
+                self.schedule_b(scheduled_node)
+            elif scheduled_node.type == 'W':
+                self.schedule_w(scheduled_node)
+            else:
+                raise ValueError(f"Unknown node type {scheduled_node.type}")
+
+        for h in self.send_handles:
+            for hh in h:
+                for hhh in hh:
+                    hhh.wait()
+
+        if not self.forward_only:
+            # Launch any remaining grad reductions
+            if self.no_sync_context is not None:
+                self.enable_grad_sync()
+
+            if self.config.finalize_model_grads_func is not None:
+                # Finalize model grads (perform full grad all-reduce / reduce-scatter for
+                # data parallelism, layernorm all-reduce for sequence parallelism).
+                self.config.finalize_model_grads_func([self.model])
+
+        return self.forward_data_store
+
+    def __call__(self, *args, **kwargs):
+        self.prepare(*args, **kwargs)
+        return self.run()
+
+zb_scheduler = ZeroBubbleScheduler()
+
+def bootstrap_p2p_communication(config):
+
+    if parallel_state.get_pipeline_model_parallel_world_size() > 1:
+        nccl_init_tensor = [torch.Tensor([torch.distributed.get_rank() + 100]).cuda() ]
+        shape = [(1,)]
+        if not parallel_state.is_pipeline_first_stage(ignore_virtual=True):
+            recv_forward(shape, config)
+        if not parallel_state.is_pipeline_last_stage(ignore_virtual=True):
+            send_forward(nccl_init_tensor, shape, config)
+            recv_backward(shape, config)
+        if not parallel_state.is_pipeline_first_stage(ignore_virtual=True):
+            send_backward(nccl_init_tensor, shape, config)
+            exit()
+        torch.distributed.barrier()
+
+def get_zero_bubble_forward_backward_func():
+    pipeline_model_parallel_size = parallel_state.get_pipeline_model_parallel_world_size()
+    assert (pipeline_model_parallel_size > 1), "zero-bubble must be used with pipelined parallelism"
+    return zb_scheduler
diff --git a/megatron-lm-musa-patch/musa_patch/zbb_light/zbpp_light.py b/megatron-lm-musa-patch/musa_patch/zbb_light/zbpp_light.py
new file mode 100644
index 000000000..c9bac0179
--- /dev/null
+++ b/megatron-lm-musa-patch/musa_patch/zbb_light/zbpp_light.py
@@ -0,0 +1,126 @@
+import os, sys
+from typing import Optional, Tuple, Union, List
+import torch
+def wrap_w_funcs(original_func):
+    from .weight_grad_store import WeightGradStore
+    def wrapped_func(total_input, grad_output, weight):
+        from megatron.training import get_args
+        if os.getenv("ENABLE_ZERO_BUBBLE", "0") == "1":
+            WeightGradStore.put((total_input, grad_output, weight), original_func)
+        else:
+            original_func(total_input, grad_output, weight)
+    return wrapped_func
+
+def wrap_w_funcs_gemm(original_func):
+    def wrapped_func(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        dtype: torch.dtype,
+        workspace: torch.Tensor,
+        layout: str = "TN",
+        **kwargs
+    ) -> Tuple[Union[torch.Tensor, None], ...]:
+
+        import functools
+        from musa_patch.zbb_light.weight_grad_store import WeightGradStore
+        from megatron.training import get_args
+        if layout == "NT" and os.getenv("ENABLE_ZERO_BUBBLE", "0") == "1":
+            WeightGradStore.put(
+                                (A, B, dtype, workspace),
+                                functools.partial(
+                                    original_func,
+                                    layout="NT",
+                                    **kwargs)
+                                )
+            return (None, None, None)
+        else:
+            return original_func(A, B, dtype, workspace, layout=layout, **kwargs)
+
+
+    return wrapped_func
+
+def wrap_w_general_gemm(original_func):
+    def wrapped_func(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        workspace: torch.Tensor,
+        layout: str = "TN",
+        **kwargs
+    ) -> Tuple[Union[torch.Tensor, None], ...]:
+
+        import functools
+        from musa_patch.zbb_light.weight_grad_store import WeightGradStore
+        from megatron.training import get_args
+        if layout == "NT" and os.getenv("ENABLE_ZERO_BUBBLE", "0") == "1":
+            from transformer_engine.pytorch.utils import clear_tensor_data
+            WeightGradStore.put(
+                                (A, B, workspace),
+                                functools.partial(
+                                    original_func,
+                                    layout="NT",
+                                    **kwargs),
+                                clear_tensor_data
+                                )
+            return (None, None, None, None)
+        else:
+            return original_func(A, B, workspace, layout=layout, **kwargs)
+
+
+    return wrapped_func
+
+def wrap_w_general_grouped_gemm(original_func):
+    def wrapped_func(
+        A: List[torch.Tensor],
+        B: List[torch.Tensor],
+        out: List[torch.Tensor],
+        out_dtype: torch.dtype,
+        workspaces: List[torch.Tensor],
+        layout: str = "TN",
+        use_bias: bool = False,
+        **kwargs
+    ) -> Tuple[Union[torch.Tensor, None], ...]:
+
+        import functools
+        from musa_patch.zbb_light.weight_grad_store import WeightGradStore
+        from megatron.training import get_args
+        if layout == "NT" and os.getenv("ENABLE_ZERO_BUBBLE", "0") == "1":
+            from transformer_engine.pytorch.utils import clear_tensor_data
+            WeightGradStore.put(
+                                (A, B, out, out_dtype, workspaces),
+                                functools.partial(
+                                    original_func,
+                                    layout="NT",
+                                    use_bias = False,
+                                    **kwargs),
+                                clear_tensor_data
+                                )
+            assert use_bias== False, "Zero-bubble doesn't support the case where bias is used."
+            return (None, [None] * len(A), None)
+        else:
+            return original_func(A, B, out, out_dtype, workspaces, use_bias = use_bias, layout=layout, **kwargs)
+
+
+    return wrapped_func
+
+
+def patch_megatron():
+    assert all([not x.startswith('megatron') for x in sys.modules.keys()]), 'Please patch zbpp before importing any megatron modules.'
+    import fused_weight_gradient_mlp_cuda
+    assert hasattr(fused_weight_gradient_mlp_cuda, 'wgrad_gemm_accum_fp32')
+    assert hasattr(fused_weight_gradient_mlp_cuda, 'wgrad_gemm_accum_fp16')
+    fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32 = wrap_w_funcs(fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32)
+    fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp16 = wrap_w_funcs(fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp16)
+
+    import transformer_engine.pytorch.cpp_extensions
+    transformer_engine.pytorch.cpp_extensions.gemm = wrap_w_funcs_gemm(transformer_engine.pytorch.cpp_extensions.gemm)
+    transformer_engine.pytorch.cpp_extensions.general_gemm = wrap_w_general_gemm(transformer_engine.pytorch.cpp_extensions.general_gemm)
+    transformer_engine.pytorch.cpp_extensions.general_grouped_gemm = wrap_w_general_grouped_gemm(transformer_engine.pytorch.cpp_extensions.general_grouped_gemm)
+
+    import megatron.core.pipeline_parallel
+    from .zb_schedule import get_zero_bubble_forward_backward_func
+    assert hasattr(megatron.core.pipeline_parallel.schedules, 'get_forward_backward_func')
+    assert hasattr(megatron.core.pipeline_parallel, 'get_forward_backward_func')
+    megatron.core.pipeline_parallel.schedules.get_forward_backward_func_origin = megatron.core.pipeline_parallel.schedules.get_forward_backward_func
+    megatron.core.pipeline_parallel.get_forward_backward_func_origin = megatron.core.pipeline_parallel.get_forward_backward_func
+    megatron.core.pipeline_parallel.schedules.get_forward_backward_func = get_zero_bubble_forward_backward_func
+    megatron.core.pipeline_parallel.get_forward_backward_func = get_zero_bubble_forward_backward_func
diff --git a/megatron/core/datasets/helpers_cpp.cpython-310-x86_64-linux-gnu.so b/megatron/core/datasets/helpers_cpp.cpython-310-x86_64-linux-gnu.so
new file mode 100755
index 0000000000000000000000000000000000000000..6125a7b497de0acd533793e830327b3c3d0f88bc
GIT binary patch
diff --git a/megatron/core/enums.py b/megatron/core/enums.py
index 46e7d3b76..c74700d65 100644
--- a/megatron/core/enums.py
+++ b/megatron/core/enums.py
@@ -8,3 +8,12 @@ class ModelType(enum.Enum):
     encoder_and_decoder = 2
     retro_encoder = 3
     retro_decoder = 4
+
+
+class Fp8Recipe(str, enum.Enum):
+    """FP8 recipe names: delayed, tensorwise, mxfp8, blockwise."""
+
+    delayed = "delayed"
+    tensorwise = "tensorwise"
+    mxfp8 = "mxfp8"
+    blockwise = "blockwise"
diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index 1d5725cc8..ab6678551 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -1344,16 +1344,30 @@ try:
 
     from transformer_engine.pytorch.permutation import (
         moe_permute,
-        moe_sort_chunks_by_index,
+        # moe_sort_chunks_by_index,
         moe_unpermute,
     )
 
     fused_permute = moe_permute
     fused_unpermute = moe_unpermute
-    fused_sort_chunks_by_index = moe_sort_chunks_by_index
+    fused_sort_chunks_by_index = None #moe_sort_chunks_by_index
 
 except ImportError:
 
     fused_permute = None
     fused_unpermute = None
     fused_sort_chunks_by_index = None
+
+try:
+
+    from transformer_engine.pytorch.cross_entropy import parallel_cross_entropy
+
+    def te_parallel_cross_entropy(logits: torch.Tensor, labels: torch.Tensor):
+        """Wrapper function for TE's Cross Entropy Loss kernel"""
+        return parallel_cross_entropy(
+            logits, labels, 0.0, False, get_tensor_model_parallel_group(check_initialized=False)
+        )
+
+except ImportError:
+
+    te_parallel_cross_entropy = None
\ No newline at end of file
diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py
index f18beada7..9a049b048 100644
--- a/megatron/core/fp8_utils.py
+++ b/megatron/core/fp8_utils.py
@@ -5,6 +5,7 @@ from typing import Tuple
 
 import torch
 
+from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import is_te_min_version
 
 # Check if Transformer Engine is installed
@@ -43,7 +44,9 @@ if HAVE_TE and is_te_min_version("2.0"):
     # TE quantization logic using quantizer API
     # Supported TE versions: 2.0+
     from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor
-
+    from megatron.core.extensions.transformer_engine import TEDelayedScaling
+    from megatron.core import parallel_state
+    from megatron.core.enums import Fp8Recipe
     def _quantize_param_fragment_impl(
         input_: torch.Tensor, *, out: torch.Tensor, param: torch.nn.Parameter
     ) -> None:
@@ -63,6 +66,129 @@ if HAVE_TE and is_te_min_version("2.0"):
         quantizer = tensor._quantizer
         return quantizer.scale, quantizer.amax
 
+    def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool = False):
+        """Return fp8 context manager.
+
+        Arguments:
+            config (TransformerConfig): Configuration object.
+            layer_no (int): *Global* layer index (including layers on other
+                pipeline-parallel ranks).
+            is_init (bool): Whether the context is fp8_model_init (True) or fp8_autocast (False).
+
+        Returns:
+            FP8 context.
+            If layer_no < 0, we return a fp8 context for all layers regardless of layer_no.
+            We return nullcontext() when: a) not using fp8 to train, b) layer_no is a layer
+            that needs to be trained in bf16.
+        """
+        num_bf16_layers_at_start = (
+            #config.num_layers_at_start_in_bf16 if config.first_last_layers_bf16 else 0
+            config.num_layers_at_start_in_bf16 if False else 0
+        )
+        num_bf16_layers_at_end = (
+            #config.num_layers_at_end_in_bf16 if config.first_last_layers_bf16 else 0
+            config.num_layers_at_end_in_bf16 if False else 0
+        )
+        # Since layer_no is a global layer index, additional checks on whether
+        # we are in the first or last pipeline-parallel rank are not needed.
+        is_first_layer = layer_no < num_bf16_layers_at_start
+        is_last_layer = layer_no >= config.num_layers - num_bf16_layers_at_end
+
+        need_fp8_context = config.fp8 if not is_init else config.fp8_param
+
+        if not need_fp8_context:
+            # bf16 training
+            fp8_context = nullcontext()
+        elif layer_no >= 0 and config.first_last_layers_bf16 and (is_first_layer or is_last_layer):
+            # fp8 training but this layer_no should be bf16
+            fp8_context = nullcontext()
+        else:
+            # fp8 training and this layer_no is in fp8
+            import transformer_engine  # To keep out TE dependency when not training in fp8
+
+            if config.fp8 == "e4m3":
+                fp8_format = transformer_engine.common.recipe.Format.E4M3
+            elif config.fp8 == "hybrid":
+                fp8_format = transformer_engine.common.recipe.Format.HYBRID
+            else:
+                raise ValueError("E4M3 and HYBRID are the only supported FP8 formats.")
+
+            # Select fp8 recipe (TE version >= 2.1.0).
+            fp8_recipe = None
+            if is_te_min_version("2.1.0"):
+                if config.fp8_recipe == Fp8Recipe.delayed:
+                    fp8_recipe = TEDelayedScaling(
+                        config=config,
+                        fp8_format=fp8_format,
+                        override_linear_precision=(False, False, not config.fp8_wgrad),
+                    )
+                elif config.fp8_recipe == Fp8Recipe.tensorwise and is_te_min_version("2.2.0.dev0"):
+                    fp8_recipe = transformer_engine.common.recipe.Float8CurrentScaling(
+                        fp8_format=fp8_format
+                    )
+                elif config.fp8_recipe == Fp8Recipe.blockwise and is_te_min_version("2.3.0.dev0"):
+                    fp8_recipe = transformer_engine.common.recipe.Float8BlockScaling(
+                        fp8_format=fp8_format
+                    )
+                elif config.fp8_recipe == Fp8Recipe.mxfp8:
+                    fp8_recipe = transformer_engine.common.recipe.MXFP8BlockScaling(
+                        fp8_format=fp8_format
+                    )
+                else:
+                    raise ValueError(
+                        "Float8CurrentScaling, MXFP8BlockScaling, Float8BlockwiseScaling and "
+                        "DelayedScaling are the only supported FP8 recipes. Please also make sure "
+                        "you are using a compatible TE version."
+                    )
+            else:
+                # Assert that the user is using delayed scaling.
+                #assert config.fp8_recipe == Fp8Recipe.delayed, (
+                #    "Please make sure to use TransformerEngine version >= 2.2.0.dev0 for "
+                #    "Float8CurrentScaling, >= 2.1.0 for MXFP8BlockScaling, and >= 2.3.0.dev0 for "
+                #    "Float8BlockScaling."
+                #)
+                fp8_recipe = TEDelayedScaling(
+                    config=config,
+                    fp8_format=fp8_format,
+                    override_linear_precision=(False, False, not config.fp8_wgrad),
+                )
+
+            fp8_group = None
+            if parallel_state.model_parallel_is_initialized():
+                fp8_group = parallel_state.get_amax_reduction_group(
+                    with_context_parallel=True, tp_only_amax_red=config.tp_only_amax_red
+                )
+
+            if not is_init:
+                fp8_context = transformer_engine.pytorch.fp8_autocast(
+                    enabled=True, fp8_recipe=fp8_recipe, fp8_group=fp8_group
+                )
+            else:
+                import inspect
+
+                context_args = {"enabled": True}
+                # Check if fp8_model_init supports setting recipe
+                if "recipe" in (
+                    inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters
+                ):
+                    context_args["recipe"] = fp8_recipe
+                # Check if fp8_model_init supports preserve_high_precision_init_val
+                if "preserve_high_precision_init_val" in (
+                    inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters
+                ):
+                    context_args["preserve_high_precision_init_val"] = True
+                fp8_context = transformer_engine.pytorch.fp8_model_init(**context_args)
+
+            # First / last layer in bf16 isn't supported with delayed scaling since it
+            # requires entering/exiting fp8 context per layer, causing incorrect amax
+            # reduction behavior.
+            #assert not (
+            #    config.first_last_layers_bf16 and isinstance(fp8_recipe, TEDelayedScaling)
+            #), "Delayed scaling does not support first / last layer in BF16."
+
+        return fp8_context
+
+
 elif HAVE_TE and is_te_min_version("1.0"):
     # TE quantization logic with fp8_meta dicts
     # Supported TE versions: 1.0 - 1.14
@@ -93,6 +219,11 @@ else:
         raise RuntimeError("Invalid Transformer Engine version for FP8 distributed optimizer")
 
 
+    def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool = False):
+        """Returns dummy fp8 context manager since TE is not available."""
+        return nullcontext()
+
+
 def quantize_param_fragment(
     input_: torch.Tensor, *, out: torch.Tensor, param: torch.nn.Parameter
 ) -> None:
diff --git a/megatron/core/inference/contexts/__init__.py b/megatron/core/inference/contexts/__init__.py
new file mode 100644
index 000000000..1b1324db7
--- /dev/null
+++ b/megatron/core/inference/contexts/__init__.py
@@ -0,0 +1,22 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import warnings
+
+from .base_context import BaseInferenceContext
+from .dynamic_chunk_allocator import ChunkAllocator
+from .static_context import StaticInferenceContext
+
+warnings.warn(
+    "The following imports from `dynamic_context.py` will be removed "
+    "in this file in `megatron-core` 0.14. The imports here result in "
+    "a cyclic import issue that causes rotary embeddings to import "
+    "from Apex rather than Transformer Engine.",
+    DeprecationWarning,
+)
+from .dynamic_context import (
+    ChunkOverflowError,
+    ContextOverflowError,
+    DynamicInferenceContext,
+    RequestOverflowError,
+    TokenOverflowError,
+)
diff --git a/megatron/core/inference/contexts/base_context.py b/megatron/core/inference/contexts/base_context.py
new file mode 100644
index 000000000..5e61a6501
--- /dev/null
+++ b/megatron/core/inference/contexts/base_context.py
@@ -0,0 +1,20 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import abc
+
+
+class BaseInferenceContext(abc.ABC):
+    """Base class for inference contexts.
+
+    Currently extended by `StaticInferenceContext` and `DynamicInferenceContext`.
+    Extend this class for any future contexts types.
+    """
+
+    @abc.abstractmethod
+    def is_static_batching(self) -> bool:
+        """Return `True` if context uses static batching."""
+        pass
+
+    def is_dynamic_batching(self) -> bool:
+        """Return `True` if context uses dynamic batching."""
+        return not self.is_static_batching()
diff --git a/megatron/core/inference/contexts/dynamic_chunk_allocator.py b/megatron/core/inference/contexts/dynamic_chunk_allocator.py
new file mode 100644
index 000000000..cbc1127dd
--- /dev/null
+++ b/megatron/core/inference/contexts/dynamic_chunk_allocator.py
@@ -0,0 +1,92 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+class ChunkAllocator:
+    """Allocator that manages chunks of memory for the KV cache.
+
+    This allocator is responsible for:
+    - Initializing a pool of chunk IDs
+    - Allocating chunks from the pool
+    - Releasing chunks back to the pool
+    - Managing the guaranteed chunk count for active requests
+
+    Args:
+        chunk_count_total (int): Total number of chunks available in the buffer.
+        gtd_chunk_count (int): Number of chunks reserved for guaranteed requests.
+    """
+
+    def __init__(self, chunk_count_total: int, gtd_chunk_count: int):
+        self.chunk_count_total = chunk_count_total
+        self.gtd_chunk_count = gtd_chunk_count
+
+        # Reserve last chunk ID as dummy chunk for decode-only inference steps
+        self.chunk_count_avail = self.chunk_count_total - 1
+        self.dummy_chunk_idx = self.chunk_count_total - 1
+
+        # Initialize chunk pool as a "stack" data structure
+        self.chunk_bag = torch.arange(
+            self.chunk_count_total, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+
+    def is_memory_available(self, num_chunks: int, safe: bool = False) -> bool:
+        """Check if memory chunks are available.
+
+        Use 'safe' to avoid all requests being blocked. A fraction of the KV cache
+        memory buffer is reserved to guarantee that a minimum number of active
+        requests can run on any given step.
+
+        Args:
+            num_chunks (int): Number of chunks to check.
+            safe (bool): Include extra space for guaranteeing ability to run
+                requests to completion.
+
+        Return:
+            (bool) Is memory available?
+        """
+        if safe:
+            return self.chunk_count_avail >= num_chunks + self.gtd_chunk_count
+        else:
+            return self.chunk_count_avail >= num_chunks
+
+    def allocate_memory_chunks(self, num_chunks: int = 1, safe: bool = False) -> Optional[Tensor]:
+        """Allocate memory chunks if available, else return None.
+
+        Args:
+            num_chunks (int): Number of chunks to allocate.
+            safe (bool): Include extra space for guaranteeing ability to run
+                requests to completion.
+
+        Return:
+            (Optional[Tensor]) Allocated chunk IDs.
+        """
+        if self.is_memory_available(num_chunks, safe):
+            self.chunk_count_avail -= num_chunks
+            return self.chunk_bag[self.chunk_count_avail : (self.chunk_count_avail + num_chunks)]
+        else:
+            return None
+
+    def release_memory_chunks(self, chunks: Tensor) -> None:
+        """Release memory chunks.
+
+        Args:
+            chunks (Tensor): Chunk IDs to release.
+
+        Return:
+            None
+        """
+        num_chunks = chunks.size(dim=0)
+        self.chunk_bag[self.chunk_count_avail : (self.chunk_count_avail + num_chunks)] = chunks
+        self.chunk_count_avail += num_chunks
+
+    def reset(self) -> None:
+        """Reset the allocator to initial state.
+
+        This resets the available chunk count to the entire memory pool
+        (except for the dummy chunk).
+        """
+        self.chunk_count_avail = self.chunk_count_total - 1
diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
new file mode 100644
index 000000000..fd327face
--- /dev/null
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -0,0 +1,1001 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import math
+import warnings
+from typing import Optional, Tuple
+
+import torch
+from packaging.version import Version as PkgVersion
+from torch import Tensor
+
+from megatron.core import parallel_state
+from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb
+from megatron.core.package_info import __version__ as mcore_version
+from megatron.core.transformer import TransformerConfig
+from megatron.core.utils import divide as core_divide
+
+from .base_context import BaseInferenceContext
+from .dynamic_chunk_allocator import ChunkAllocator
+
+
+class ContextOverflowError(Exception):
+    '''Base exception for when a new request would not fit.'''
+
+    pass
+
+
+class RequestOverflowError(ContextOverflowError):
+    '''Adding request would overflow max request count.'''
+
+    pass
+
+
+class TokenOverflowError(ContextOverflowError):
+    '''Adding request would overflow max token count.'''
+
+    pass
+
+
+class ChunkOverflowError(ContextOverflowError):
+    '''Adding request would overflow available memory chunks.'''
+
+    pass
+
+
+# pylint: disable=line-too-long
+class DynamicInferenceContext(BaseInferenceContext):
+    """Inference context that is passed to the main model in order
+    to efficiently calculate and store the KV cache during inference.
+
+    The dynamic inference context manages both: 1) in-flight batching, and 2) a
+    memory buffer for the chunked KV cache. For in-flight batching, requests of
+    arbitrary sequence length may be added, paused, or removed from the context
+    at any step. The only constraint is the maximum number of requests or tokens
+    that the context is defined to support. For the chunked KV cache, a memory
+    buffer is allocated up front (size `buffer_size_gb`), that is divided into
+    chunks and dynamically assigned to requests. At any given step, any unassigned
+    chunks equate to unused space.
+
+    Additionally, a fraction of the memory buffer (`gtd_request_fraction`, i.e.,
+    the 'guaranteed' request fraction) is reserved for guaranteeing that a
+    minimum number of active requests may continue to generate tokens on any step.
+    The reason for this is that the context manages two pools of requests: 1)
+    active requests, and 2) paused requests. Paused requests are requests where
+    insufficient memory chunks remain for future assignment, and these requests
+    are set aside until enough memory chunks are available. Active requests are
+    requests that have sufficient memory chunks to proceed with their generations.
+
+    The situation can arise where all requests eventually become paused due to all
+    memory chunks being assigned. In this case, there are no active requests and
+    thus no progress can be made. To handle this case, a fraction of the memory
+    buffer is reserved that only allows active requests, and no paused requests.
+    This fraction must be carefully tuned, as it can have an order of magnitude
+    impact on overall latency.
+
+    Args:
+        params_dtype (torch.dtype): Dtype used for KV cache.
+        num_layers (int): Number of layers.
+        kv_channels (int): Hidden dimension per attention head.
+        num_attention_heads (int): Number of attention heads.
+        max_sequence_length (int): Max possible sequence length (prompt + output)
+            that will occur.
+        buffer_size_gb (float): Total buffer size (GB), shared by main and
+            fallback contexts.
+        chunk_size_tokens (int): Size of KV cache chunk size.
+        buffer_guaranteed_fraction (float): Fraction of the memory buffer that is
+            reserved to guarantee that one or more active requests are able to
+            run to completion. Without reserving this memory, paused requests are
+            able to fill the memory buffer and block execution of any requests.
+        buffer_overflow_factor (Optional[float]): Scaling factor over the buffer
+            size for auto computing `max_requests` and `max_tokens`. This scaling
+            factor is used for fitting more requests and tokens in the memory
+            buffer than it can safely hold, which in turn increases throughput.
+        max_requests_override (Optional[int]): If set, overrides value computed
+            from `buffer_overflow_factor`.
+        max_tokens_override (Optional[int]): If set, overrides value computed
+            from `buffer_overflow_factor`.
+    """
+
+    def __init__(
+        self,
+        *,
+        params_dtype: torch.dtype,
+        num_layers: int,
+        kv_channels: int,
+        num_attention_heads: int,
+        max_sequence_length: int,
+        buffer_size_gb: float,
+        buffer_guaranteed_fraction: float,
+        chunk_size_tokens: int = 256,
+        buffer_overflow_factor: Optional[float] = None,
+        max_requests_override: Optional[int] = None,
+        max_tokens_override: Optional[int] = None,
+        tensor_model_parallel_size: Optional[int] = None,
+    ):
+
+        super().__init__()
+        # Per partition num heads and hidden size.
+        projection_size = kv_channels * num_attention_heads
+        if tensor_model_parallel_size is None:
+            tp_size = parallel_state.get_tensor_model_parallel_world_size()
+        else:
+            tp_size = tensor_model_parallel_size
+        hidden_size_per_attention_head = core_divide(projection_size, num_attention_heads)
+        num_attention_heads_per_partition = core_divide(num_attention_heads, tp_size)
+
+        # Chunk size tokens, bytes.
+        dtype_size_bytes = params_dtype.itemsize
+        self.chunk_size_tokens = chunk_size_tokens
+        self.chunk_size_bytes = (
+            dtype_size_bytes
+            * 2  # key, value
+            * num_layers
+            * self.chunk_size_tokens
+            * num_attention_heads_per_partition
+            * hidden_size_per_attention_head
+        )
+
+        # Adjust buffer to be a multiple of chunk size.
+        buffer_size_bytes = int(buffer_size_gb * 1024**3)
+        buffer_size_bytes_rem = buffer_size_bytes % self.chunk_size_bytes
+        buffer_size_bytes = buffer_size_bytes - buffer_size_bytes_rem
+
+        # Compute max_requets, max_tokens from buffer size and overflow factor.
+        def bytes_to_max_requests_and_tokens(n_bytes):
+            n_tokens = n_bytes / self.chunk_size_bytes * self.chunk_size_tokens
+            n_requests = n_tokens / max_sequence_length
+            return int(n_requests), int(n_tokens)
+
+        self.max_requests, self.max_tokens = bytes_to_max_requests_and_tokens(buffer_size_bytes)
+
+        if buffer_overflow_factor is not None:
+            self.max_requests = self.round_up_requests(
+                int(self.max_requests * buffer_overflow_factor)
+            )
+            self.max_tokens = self.round_up_tokens(
+                int(self.max_tokens * buffer_overflow_factor / 50.0)
+            )
+
+        if max_requests_override is not None:
+            self.max_requests = self.round_up_requests(max_requests_override)
+
+        if max_tokens_override is not None:
+            self.max_tokens = self.round_up_tokens(max_tokens_override)
+
+        self.max_requests = min(self.max_requests, self.max_tokens)  # e.g., decode only.
+
+        # Initialize context state.
+        self.params_dtype = params_dtype
+        self.num_layers = num_layers
+        self.max_sequence_length = max_sequence_length
+
+        self.total_request_count = 0
+        self.active_token_count = 0
+        self.paused_request_count = 0
+        self.padded_active_token_count = None
+        self.padded_active_sample_count = None
+        self.paused_tokens = None
+
+        # Per-request state.
+        self.request_ids = torch.full(
+            (self.max_requests,), -1, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+        # request_query_lengths is the input prompt tokens length during prefill phase (1st step) and then 1 for the decode phase (i.e During generation)
+        self.request_query_lengths = torch.empty_like(self.request_ids)
+        # request_output_lengths is len(input_prompt_tokens) + num_tokens_to_generate
+        self.request_output_lengths = torch.empty_like(self.request_ids)
+        # request_kv_length_offsets is the same as query length during prefill phase (1st step) and then 1 for the decode phase (i.e During generation)
+        self.request_kv_length_offsets = torch.empty_like(self.request_ids)
+        self.request_kv_chunk_counts = torch.empty_like(self.request_ids)
+        self.request_last_kv_chunk_id = torch.empty_like(self.request_ids)
+        # request_last_kv_chunk_offset represents number of tokens in the last kv chunk
+        self.request_last_kv_chunk_offset = torch.empty_like(self.request_ids)
+
+        # Per-token state.
+        self.token_to_input_ids = torch.full(
+            (self.max_tokens,), 0, dtype=torch.long, device=torch.cuda.current_device()
+        )
+        self.token_to_pos_ids = torch.full_like(self.token_to_input_ids, 0)
+        self.token_to_request_idx = torch.empty_like(self.token_to_input_ids)
+        self.token_to_chunk_idx = torch.empty_like(self.token_to_input_ids)
+        # i.e For a set of tokens A B C D E F ..  and chunk_size 4:
+        # token_to_position_in_request is  [0, 1, 2, 3, 4, 5]
+        # token_to_local_position_within_kv_chunk is [0 , 1, 2, 3, 0, 1, 2]
+        self.token_to_position_in_request = torch.empty_like(self.token_to_input_ids)
+        self.token_to_local_position_within_kv_chunk = torch.empty_like(self.token_to_input_ids)
+
+        # Calculate the total number of chunks available in the buffer
+        chunk_count_total = buffer_size_bytes // self.chunk_size_bytes
+
+        # Memory buffer.
+        self.memory_buffer = torch.full(
+            (
+                2,  # key and value
+                self.num_layers,
+                chunk_count_total,
+                self.chunk_size_tokens,
+                num_attention_heads_per_partition,
+                hidden_size_per_attention_head,
+            ),
+            -1,
+            dtype=self.params_dtype,
+            device=torch.cuda.current_device(),
+        )
+
+        # Chunk ids.
+        self.max_kv_chunk_count = math.ceil(self.max_sequence_length / self.chunk_size_tokens)
+        self.request_to_kv_chunk_ids = torch.full(
+            (self.max_requests, self.max_kv_chunk_count),
+            -1,
+            dtype=torch.int,
+            device=torch.cuda.current_device(),
+        )
+
+        # `*_decode_only` tensors are for use with cuda graphs to maintain
+        # consistent input shapes, which is required to use cuda graphs. Cuda
+        # graphs are used only during decode-only steps (i.e., no requests are in
+        # the prefill phases). During these decode-only steps, the `*_decode_only`
+        # tensors are used, otherwise their same-name but un-suffixed
+        # corresponding tensors are used.
+        # TODO: @lmcafee, only use `_decode_only` tensors when both of the
+        # following conditions are met: 1) decode-only step, and 2) cuda graphs
+        # are enabled.
+
+        self.query_seq_lengths_decode_only = torch.full(
+            (self.max_requests,), 0, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+        self.cu_query_seq_lengths_decode_only = torch.full(
+            (self.max_requests + 1,), 0, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+        self.kv_seq_lengths_decode_only = torch.full(
+            (self.max_requests,), 0, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+        self.cu_kv_seq_lengths_decode_only = torch.full(
+            (self.max_requests + 1,), 0, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+
+        self.kv_memory_decode_only = torch.full(
+            (self.max_requests, self.max_kv_chunk_count),
+            0,
+            dtype=torch.int,
+            device=torch.cuda.current_device(),
+        )
+
+        # Guaranteed active requests.
+        # * See details in the class docstring above. `gtd_request_fraction` is
+        #   the fraction of the memory buffer that is reserved for guaranteeing
+        #   that some number of active requests can always proceed with their
+        #   generations. The number of bytes defined by `gtd_request_fraction *
+        #   buffer_size_gb` is converted to a number of requests that this
+        #   reserved space can handle (`gtd_request_count`), and rounded to be an
+        #   exact multiple of `max_sequence_length`. This is then converted into
+        #   the number of reserved chunks (`gtd_chunk_count`) and bytes
+        #   (`gtd_byte_count`).
+        # Chunk ids.
+        self.max_kv_chunk_count = math.ceil(self.max_sequence_length / self.chunk_size_tokens)
+        gtd_byte_count = buffer_guaranteed_fraction * buffer_size_bytes
+        gtd_request_count, _ = bytes_to_max_requests_and_tokens(gtd_byte_count)
+        if buffer_guaranteed_fraction > 0:
+            gtd_request_count = max(1, gtd_request_count)
+        gtd_request_count = self.round_up_requests(min(gtd_request_count, self.max_requests))
+        gtd_chunk_count = gtd_request_count * self.max_kv_chunk_count
+        assert (
+            gtd_request_count <= self.max_requests
+        ), "gtd_request_count (%d) > max_requests (%d)." % (gtd_request_count, self.max_requests)
+        self.gtd_request_count = gtd_request_count
+        self.gtd_chunk_count = gtd_chunk_count
+
+        # Initialize chunk allocator
+        self.chunk_allocator = ChunkAllocator(
+            chunk_count_total=chunk_count_total, gtd_chunk_count=self.gtd_chunk_count
+        )
+
+        # Store the dummy chunk idx reference for convenience
+        self.dummy_chunk_idx = self.chunk_allocator.dummy_chunk_idx
+        # Reset attention state.
+        self.reset_attention_state()
+
+    TOKEN_ROUNDER = 64
+    REQUEST_ROUNDER = 4
+
+    @classmethod
+    def round_up_tokens(cls, value):
+        """Round up to nearest multiple of `TOKEN_ROUNDER` (above)."""
+        if PkgVersion(mcore_version) < PkgVersion("0.13"):
+            return cls.round_up(value)
+        return cls.TOKEN_ROUNDER * int(math.ceil(int(value) / cls.TOKEN_ROUNDER))
+
+    @classmethod
+    def round_up_requests(cls, value):
+        """Round up to nearest multiple of `REQUEST_ROUNDER` (above)."""
+        if PkgVersion(mcore_version) < PkgVersion("0.13"):
+            return cls.round_up(value)
+        return cls.REQUEST_ROUNDER * int(math.ceil(int(value) / cls.REQUEST_ROUNDER))
+
+    @classmethod
+    def round_up(cls, value):
+        """Deprecated in favor of round_up_tokens and round_up_requests."""
+        warnings.warn(
+            "`round_up` is deprecated in favor of `round_up_tokens` or `round_up_requests` "
+            "and will be removed in `megatron-core` 0.14."
+        )
+        ROUNDER = getattr(cls, "ROUNDER", 64)
+        return ROUNDER * int(math.ceil(int(value) / ROUNDER))
+
+    def is_static_batching(self) -> bool:
+        """Is static batching? False."""
+        return False
+
+    def is_decode_only(self) -> bool:
+        """Test if all active requests are in decode phase.
+
+        For a request in prefill phase active_tokens = query length
+        Once the request moves to decode phase active tokens is 1 for that request. So if all active requests are in decode phase, they will be equal to active token count.
+        """
+        total_active_requests = self.total_request_count - self.paused_request_count
+        return total_active_requests == self.active_token_count
+
+    def has_unfinished_requests(self) -> bool:
+        """Test if any requests remain."""
+        return self.total_request_count > 0
+
+    def cu_query_lengths(self) -> Tensor:
+        """Cumulative query sequence lengths."""
+        return self.cu_query_seq_lengths, self.max_seqlen_q
+
+    def cu_kv_lengths(self) -> Tensor:
+        """Cumulative key/value sequence lengths."""
+        return (
+            self.cu_kv_seq_lengths,
+            self.kv_seq_lengths,
+            self.kv_seq_lengths_decode_only,
+            self.max_seqlen_k,
+        )
+
+    def get_active_sequence_lengths(self) -> Tensor:
+        """Total sequence length (query + key) for active requests."""
+        lengths = self.request_kv_length_offsets + self.request_query_lengths
+        lengths = lengths[self.paused_request_count : self.total_request_count]
+        return lengths
+
+    def get_max_sequence_lengths(self) -> Tensor:
+        """Maximum sequence length for active requests."""
+        return self.request_output_lengths[self.paused_request_count : self.total_request_count]
+
+    def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
+        """Append to KV cache.
+
+        Args:
+            layer_number (int): Layer number.
+            key (Tensor): Key tensor.
+            value (Tensor): Value tensor.
+        """
+
+        chunk_idx = self.token_to_chunk_idx[: self.padded_active_token_count]
+        local_kv_seq_idx = self.token_to_local_position_within_kv_chunk[
+            : self.padded_active_token_count
+        ]
+        assert key.size(1) == 1 and value.size(1) == 1
+        key = key.squeeze(1)
+        value = value.squeeze(1)
+
+        self.memory_buffer[0, layer_number - 1, chunk_idx, local_kv_seq_idx] = key[
+            : self.padded_active_token_count
+        ]
+        self.memory_buffer[1, layer_number - 1, chunk_idx, local_kv_seq_idx] = value[
+            : self.padded_active_token_count
+        ]
+
+    def key_value_cache(self, layer_number: int) -> Tuple[Tensor, Tensor]:
+        """Read from KV cache.
+
+        Args:
+            layer_number (int): Layer number.
+
+        Return:
+            (Tuple[Tensor, Tensor]) The key and value pointer tensors that point
+            to chunks within the chunked memory buffer.
+        """
+        return (
+            self.memory_buffer[0, layer_number - 1],
+            self.memory_buffer[1, layer_number - 1],
+            self.block_table,
+        )
+
+    def apply_rotary_emb_query(
+        self,
+        query: Tensor,
+        query_emb: Tensor,
+        config: TransformerConfig,
+        cu_seqlens_q: Tensor,
+        cp_group: torch.distributed.ProcessGroup,
+    ) -> Tensor:
+        """Apply rotary embedding to query tensor.
+
+        Args:
+            query (Tensor): Query tensor.
+            query_emb (Tensor): Query rotary embeddings.
+            config (TransformerConfig): Transformer config.
+            cu_seqlens_q (Tensor): Cumulative sequence lengths.
+            cp_group (torch.distributed.ProcessGroup): Process group for context parallel.
+
+        Return:
+            (Tensor) Query tensor after applying rotary embeddings.
+        """
+        n = self.padded_active_token_count
+        query_seq_idx = self.token_to_pos_ids[:n]
+        query_emb = query_emb[query_seq_idx]
+        query[:n] = apply_rotary_pos_emb(
+            t=query[:n],
+            freqs=query_emb[:n],
+            config=config,
+            cu_seqlens=cu_seqlens_q,
+            cp_group=cp_group,
+        )
+        return query
+
+    def apply_rotary_emb_key(
+        self,
+        key: Tensor,
+        key_emb: Tensor,
+        config: TransformerConfig,
+        cp_group: torch.distributed.ProcessGroup,
+    ) -> Tensor:
+        """Apply rotary embedding to key tensor.
+
+        Args:
+            key (Tensor): Key tensor.
+            key_emb (Tensor): Key rotary embeddings.
+            config (TransformerConfig): Transformer config.
+            cp_group (torch.distributed.ProcessGroup): Process group for context parallel.
+
+        Return:
+            (Tensor) Key tensor after applying rotary embeddings.
+        """
+        n = self.padded_active_token_count
+        key_seq_idx = self.token_to_position_in_request[:n]
+        key_emb = key_emb[key_seq_idx]
+        if self.is_decode_only():
+            assert key.shape[0] == n == self.max_requests
+            key = apply_rotary_pos_emb(
+                t=key[:n], freqs=key_emb[:n], config=config, cp_group=cp_group
+            )
+        else:
+            key[:n] = apply_rotary_pos_emb(
+                t=key[:n], freqs=key_emb[:n], config=config, cp_group=cp_group
+            )
+        return key
+
+    def reset_attention_state(self) -> None:
+        """Reset state used within attention, after each step."""
+        self.max_seqlen_q = None
+        self.max_seqlen_k = None
+        self.cu_query_seq_lengths = None
+        self.cu_query_seq_lengths_decode_only.fill_(0)
+        self.query_seq_lengths_decode_only.fill_(0)
+        self.cu_kv_seq_lengths = None
+        self.cu_kv_seq_lengths_decode_only.fill_(0)
+        self.kv_seq_lengths_decode_only.fill_(0)
+        self.kv_memory_decode_only.fill_(0)
+        self.block_table = None
+
+    def initialize_attention_state(self) -> None:
+        """Initialize attention state so that every layer can use it"""
+
+        self.padded_active_token_count = (
+            self.max_requests
+            if self.is_decode_only()
+            else self.round_up_tokens(self.active_token_count)
+        )
+        self.padded_active_sample_count = (
+            self.max_requests
+            if self.is_decode_only()
+            else (self.total_request_count - self.paused_request_count)
+        )
+        self.token_to_chunk_idx[self.active_token_count : self.padded_active_token_count] = (
+            self.dummy_chunk_idx
+        )
+        self.token_to_local_position_within_kv_chunk[
+            self.active_token_count : self.padded_active_token_count
+        ] = 0
+        self.token_to_position_in_request[
+            self.active_token_count : self.padded_active_token_count
+        ] = 0
+
+        query_lengths = self.request_query_lengths[
+            self.paused_request_count : self.total_request_count
+        ]
+        if self.is_decode_only():
+            self.query_seq_lengths_decode_only[
+                0 : self.total_request_count - self.paused_request_count
+            ] = query_lengths
+            cu_query_lengths_decode_only = torch.cumsum(self.query_seq_lengths_decode_only, dim=0)
+            self.cu_query_seq_lengths_decode_only[1:] = cu_query_lengths_decode_only
+            self.cu_query_seq_lengths = self.cu_query_seq_lengths_decode_only
+            self.max_seqlen_q = 1
+        else:
+            cu_query_lengths = torch.cumsum(query_lengths, dim=0)
+            self.cu_query_seq_lengths = torch.full(
+                (self.total_request_count - self.paused_request_count + 1,),
+                0,
+                dtype=torch.int32,
+                device=torch.cuda.current_device(),
+            )
+            self.cu_query_seq_lengths[1:] = cu_query_lengths
+            self.max_seqlen_q = query_lengths.max().item()
+
+        kv_seq_lengths = self.request_kv_length_offsets + self.request_query_lengths
+        self.kv_seq_lengths = kv_seq_lengths[self.paused_request_count : self.total_request_count]
+        if self.is_decode_only():
+            self.kv_seq_lengths_decode_only[
+                0 : self.total_request_count - self.paused_request_count
+            ] = self.kv_seq_lengths
+            cu_kv_lengths_decode_only = torch.cumsum(self.kv_seq_lengths_decode_only, dim=0)
+            self.cu_kv_seq_lengths_decode_only[1:] = cu_kv_lengths_decode_only
+            self.cu_kv_seq_lengths = self.cu_kv_seq_lengths_decode_only
+            self.max_seqlen_k = self.max_sequence_length
+        else:
+            self.cu_kv_seq_lengths = torch.full(
+                (self.total_request_count - self.paused_request_count + 1,),
+                0,
+                dtype=torch.int32,
+                device=torch.cuda.current_device(),
+            )
+            self.cu_kv_seq_lengths[1:] = torch.cumsum(self.kv_seq_lengths, dim=0)
+            self.max_seqlen_k = self.kv_seq_lengths.max().item()
+
+        kv_memory = self.request_to_kv_chunk_ids[
+            self.paused_request_count : self.total_request_count
+        ]
+        if self.is_decode_only():
+            self.kv_memory_decode_only[0 : self.total_request_count - self.paused_request_count] = (
+                kv_memory
+            )
+            self.block_table = self.kv_memory_decode_only
+        else:
+            self.block_table = self.request_to_kv_chunk_ids[
+                self.paused_request_count : self.total_request_count
+            ]
+
+    def reset(self) -> None:
+        """Reset entire context.
+
+        This method does:
+        - Reset active/paused request/token counts to zero.
+        - Reset available chunks to entire memory.
+        - Reset other tensors to zeros (unncessary, just or sanity checking).
+
+        This method is useful after cuda graph warmup iterations, where the
+        context's memory buffer is referenced by the cuda graph system and
+        cannot be deallocated.
+        """
+
+        # Reset request/token counts.
+        self.total_request_count = 0
+        self.active_token_count = 0
+        self.paused_request_count = 0
+        self.padded_active_token_count = 0
+        self.padded_active_sample_count = 0
+        self.paused_tokens = None
+
+        # Reset request indexes.
+        self.request_ids.fill_(-1)
+        self.request_query_lengths.fill_(0)
+        self.request_output_lengths.fill_(0)
+        self.request_kv_length_offsets.fill_(0)
+        self.request_kv_chunk_counts.fill_(0)
+        self.request_last_kv_chunk_id.fill_(-1)
+        self.request_last_kv_chunk_offset.fill_(0)
+        self.request_to_kv_chunk_ids.fill_(-1)
+
+        # Reset token indexes.
+        self.token_to_input_ids.fill_(0)
+        self.token_to_pos_ids.fill_(0)
+        self.token_to_request_idx.fill_(-1)
+        self.token_to_position_in_request.fill_(0)
+        self.token_to_chunk_idx.fill_(-1)
+        self.token_to_local_position_within_kv_chunk.fill_(0)
+
+        # Reset available chunk count.
+        self.reset_attention_state()
+        self.chunk_allocator.reset()
+        self.request_to_kv_chunk_ids.fill_(-1)
+
+    def current_input_ids(self) -> Tensor:
+        """Flattened input IDs for forward pass.
+
+        Return:
+            (Tensor) Flattened active input IDs.
+        """
+        return self.token_to_input_ids[: self.padded_active_token_count].unsqueeze(0)
+
+    def current_position_ids(self) -> Tensor:
+        """Flattened position IDs for forward pass.
+
+        Return:
+            (Tensor) Flattened active position IDs.
+        """
+        return self.token_to_pos_ids[: self.padded_active_token_count].unsqueeze(0)
+
+    def last_token_logits(self, logits: Tensor) -> Tensor:
+        """Last tokens of logits.
+
+        Args:
+            logits (Tensor): Output logits of forward pass.
+
+        Return:
+            (Tensor) Last token logits.
+        """
+
+        # todo: @lmcafee, remove these asserts?
+        assert logits.size(0) == 1
+        assert logits.size(1) == self.padded_active_token_count, (
+            f"logits.size(1) ({tuple(logits.shape)}) != "
+            f"padded_active_token_count ({self.padded_active_token_count})."
+        )
+
+        # Last token logits.
+        logits = logits.squeeze(0)
+        last_token_idxs = (
+            torch.cumsum(
+                self.request_query_lengths[self.paused_request_count : self.total_request_count],
+                dim=0,
+            )
+            - 1
+        )
+        last_token_logits = logits[last_token_idxs, :]
+
+        return last_token_logits
+
+    def add_request(
+        self, request_id: int, tokens: torch.Tensor, num_tokens_to_generate: Optional[int] = None
+    ) -> None:
+        """Add request to context.
+
+        After a request is added, it will first do one prefill step, followed by
+        an arbitrary number of decode steps.
+
+        A request will failed to be added if one of the following is true:
+        - Adding the request would overflow the max token count.
+        - Adding the request would overflow the max request count.
+        - Adding the request would overflow memory.
+
+        todo: @lmcafee, cache non-added requests until there is space, for better
+        user experience.
+
+        Args:
+            request_id (int): Unique ID of request.
+            tokens (torch.Tensor): Token IDs of request prompt.
+            num_tokens_to_generate (int): Number of tokens to generate for the request.
+
+        Return:
+            None
+        """
+
+        # `context_length` here is the equal to prompt length, and does not
+        # include output length.
+        context_length = len(tokens)
+
+        # Test for token and request overflow.
+        # TODO : Should move this into some waiting queue
+        if self.active_token_count + context_length > self.max_tokens:
+            raise TokenOverflowError()
+        if self.total_request_count >= self.max_requests:
+            raise RequestOverflowError()
+
+        # Preallocate chunks.
+        num_chunks_needed = math.ceil(context_length / self.chunk_size_tokens)
+        new_chunk_ids = self.chunk_allocator.allocate_memory_chunks(num_chunks_needed, safe=True)
+        if new_chunk_ids is None:
+            raise ChunkOverflowError()
+
+        if num_tokens_to_generate is None:
+            num_tokens_to_generate = self.max_sequence_length - context_length
+        elif context_length + num_tokens_to_generate > self.max_sequence_length:
+            raise TokenOverflowError()
+
+        # Update request state.
+        self.request_ids[self.total_request_count] = request_id
+        self.request_query_lengths[self.total_request_count] = context_length
+        self.request_output_lengths[self.total_request_count] = (
+            context_length + num_tokens_to_generate
+        )
+        self.request_kv_length_offsets[self.total_request_count] = 0
+        self.request_to_kv_chunk_ids[self.total_request_count][:num_chunks_needed] = new_chunk_ids
+        self.request_kv_chunk_counts[self.total_request_count] = num_chunks_needed
+        self.request_last_kv_chunk_id[self.total_request_count] = new_chunk_ids[-1]
+        self.request_last_kv_chunk_offset[self.total_request_count] = (
+            context_length - 1
+        ) % self.chunk_size_tokens
+
+        # Update token state.
+        arange_context_length = torch.arange(context_length, device=torch.cuda.current_device())
+
+        self.token_to_pos_ids[
+            self.active_token_count : (self.active_token_count + context_length)
+        ] = arange_context_length
+        self.token_to_input_ids[
+            self.active_token_count : (self.active_token_count + context_length)
+        ] = tokens
+
+        self.token_to_request_idx[
+            self.active_token_count : (self.active_token_count + context_length)
+        ] = self.total_request_count
+        self.token_to_position_in_request[
+            self.active_token_count : (self.active_token_count + context_length)
+        ] = arange_context_length
+        self.token_to_chunk_idx[
+            self.active_token_count : (self.active_token_count + context_length)
+        ] = new_chunk_ids[arange_context_length // self.chunk_size_tokens]
+        self.token_to_local_position_within_kv_chunk[
+            self.active_token_count : (self.active_token_count + context_length)
+        ] = (arange_context_length % self.chunk_size_tokens)
+
+        # Increment request and token counts.
+        self.total_request_count += 1
+        self.active_token_count += context_length
+
+    def _swap_book_keeping_tensors(self, src_idxs, dst_idxs, next_tokens):
+        """
+        Swaps all the relevent booking tensors with src idxs to dst idxs
+        """
+        self.request_kv_length_offsets[dst_idxs] = self.request_kv_length_offsets[src_idxs]
+        self.request_query_lengths[dst_idxs] = self.request_query_lengths[src_idxs]
+        self.request_output_lengths[dst_idxs] = self.request_output_lengths[src_idxs]
+        self.request_ids[dst_idxs] = self.request_ids[src_idxs]
+        next_tokens[dst_idxs] = next_tokens[src_idxs]
+
+        self.request_to_kv_chunk_ids[dst_idxs] = self.request_to_kv_chunk_ids[src_idxs]
+        self.request_kv_chunk_counts[dst_idxs] = self.request_kv_chunk_counts[src_idxs]
+        self.request_last_kv_chunk_id[dst_idxs] = self.request_last_kv_chunk_id[src_idxs]
+        self.request_last_kv_chunk_offset[dst_idxs] = self.request_last_kv_chunk_offset[src_idxs]
+
+    # TODO: see if we can compile this function
+    def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> None:
+        """Update context state after calling engine.step().
+
+        This method is responsible for:
+        - Update prefill requests to decode requests.
+        - Persist decode requests as decode requests.
+        - Terminate requests by length or termination id.
+
+        *Note*: All bookkeeping tensors (i.e., `self.request_*`) are laid out
+        contiguously, with a conceptual division between paused requests on the
+        'left' (or, lower indices) and active requests in the 'middle' (or, middle
+        indices) and completed requests on the 'right' (or, higher indices). The integers
+        `paused_request_count` and `total_request_count`  are used to track the boundaries
+        between these request groups.
+        - 0:paused_request_count -> paused requests
+        - paused_request_count:total_request_count -> active requests
+        - total_request_count:max_requests -> completed requests are moved here.
+        The reason for maintaining contiguous tensors rather than multiple
+        smaller (e.g., per-group or per-request) tensors is for both 1) speed
+        (avoid unnecessary tensor allocations), and 2) compatibility with the
+        Flash Attention kernels, which packed contiguous tensors.
+
+        The following happens in this code :
+        1. The active token mask tells us which requests are still active and which are completed
+        2. If no paused requests are present and no active requests we release all memory and reset.
+        3. Concatenate the paused tokens to the active tokens
+        4. For the finished requests we release memory chunks and move them to the right
+        5. We identify requests that require a new chunk and add them to the paused requests (i.e move them left)
+        6. We determine how many requests we can resume and resume them
+        7. We make changes to the request book keeping tesnsors and setup the tokens for next iteration
+        8. We resume those requests by assigning chunks and updating bookkeeping tensors
+        9. We make relevant changes to the token bookkeeping tensors
+
+        Args:
+            active_requests_mask (Tensor): 1D Mask tensor marking active requests.
+            new_tokens (Tensor): Newly sampled tokens, with one token per active request.
+
+        Return:
+            None
+        """
+        # 1. The active token mask tells us which requests are still active and which are completed
+        # active_request_count -> This corresponds to requests that have not reached EOD or max length
+        # finished_request_count are requests that have reached the termination criterion
+        active_request_count = (active_requests_mask == 1).sum().item()
+        finished_request_count = (active_requests_mask == 0).sum().item()
+        assert (
+            active_request_count + finished_request_count + self.paused_request_count
+            == self.total_request_count
+        )
+
+        # Reset attention state.
+        self.reset_attention_state()
+
+        # 2. If no paused requests are present and no active requests we release memory and reset.
+        if active_request_count + self.paused_request_count == 0:
+            if finished_request_count > 0:
+                finished_idxs = (
+                    torch.nonzero(active_requests_mask == 0, as_tuple=True)[0]
+                    + self.paused_request_count
+                )
+                kv_chunks_assigned = self.request_to_kv_chunk_ids[finished_idxs]
+                non_zero_values_in_kv_memory = kv_chunks_assigned[kv_chunks_assigned != -1]
+                self.chunk_allocator.release_memory_chunks(non_zero_values_in_kv_memory)
+
+            # Reset request/token counts.
+            self.request_to_kv_chunk_ids.fill_(-1)
+            self.total_request_count = 0
+            self.active_token_count = 0
+            return
+
+        # 3. Concatenate the paused tokens to the active tokens if present.
+        if self.paused_request_count != 0:
+            assert self.paused_tokens is not None
+            next_tokens = torch.cat((self.paused_tokens, new_tokens))
+        else:
+            next_tokens = new_tokens
+
+        # 4. For the finished requests we release memory chunks and move them to the right:-
+        #       a) Release all their memory
+        #       b) Swap them to the right, so that we have this order [Paused, Active, Finished]
+        if finished_request_count > 0:
+            finished_idxs = (
+                torch.nonzero(active_requests_mask == 0, as_tuple=True)[0]
+                + self.paused_request_count
+            )
+            kv_chunks_asigned = self.request_to_kv_chunk_ids[finished_idxs]
+            non_zero_values_in_kv_memory = kv_chunks_asigned[kv_chunks_asigned != -1]
+            self.chunk_allocator.release_memory_chunks(non_zero_values_in_kv_memory)
+
+            if active_request_count > 0:
+                finished_idxs_on_left = (
+                    torch.nonzero(active_requests_mask[:active_request_count] == 0, as_tuple=True)[
+                        0
+                    ]
+                    + self.paused_request_count
+                )
+                active_idxs_on_right = (
+                    torch.nonzero(active_requests_mask[active_request_count:], as_tuple=True)[0]
+                    + active_request_count
+                    + self.paused_request_count
+                )
+
+                self._swap_book_keeping_tensors(
+                    src_idxs=active_idxs_on_right,
+                    dst_idxs=finished_idxs_on_left,
+                    next_tokens=next_tokens,
+                )
+
+        # 5. We identify requests that require a new chunk and add them to the paused requests (i.e move them left) :-
+        #       a) Put requests that have filled their current chunk and  require a new one in a pause state temporarily
+        #       b) Move the paused requests to the left, and active requets to the right
+        #       c) Update the paused request count and active_request_count appropriately
+        if active_request_count > 0:
+            num_tokens_in_last_chunk = self.request_last_kv_chunk_offset[
+                self.paused_request_count : (active_request_count + self.paused_request_count)
+            ]
+            active_requests_requiring_new_chunk = (
+                num_tokens_in_last_chunk == self.chunk_size_tokens - 1
+            ).byte()
+            active_requests_requiring_new_chunk_count = (
+                (active_requests_requiring_new_chunk == 1).sum().item()
+            )
+
+            # Swap unfinished active requests on the left side with paused requests on the right side
+            # NOTE : We add paused request count because we concatenate
+            # paused tokens to the left at the beginning of update requests
+            if (
+                active_requests_requiring_new_chunk_count > 0
+                and active_requests_requiring_new_chunk_count != active_request_count
+            ):
+                active_request_ids_on_left = (
+                    torch.nonzero(
+                        active_requests_requiring_new_chunk[
+                            :active_requests_requiring_new_chunk_count
+                        ]
+                        == 0,
+                        as_tuple=True,
+                    )[0]
+                    + self.paused_request_count
+                )
+                paused_requests_idxs_on_right = (
+                    torch.nonzero(
+                        active_requests_requiring_new_chunk[
+                            active_requests_requiring_new_chunk_count:
+                        ],
+                        as_tuple=True,
+                    )[0]
+                    + active_requests_requiring_new_chunk_count
+                    + self.paused_request_count
+                )
+                dst_idxs = torch.cat((active_request_ids_on_left, paused_requests_idxs_on_right))
+                src_idxs = torch.cat((paused_requests_idxs_on_right, active_request_ids_on_left))
+                self._swap_book_keeping_tensors(
+                    src_idxs=src_idxs, dst_idxs=dst_idxs, next_tokens=next_tokens
+                )
+
+            self.paused_request_count += active_requests_requiring_new_chunk_count
+            active_request_count -= active_requests_requiring_new_chunk_count
+
+        # 6. Now that we have the requests in following order [Paused, Active, Finished]
+        # We determine how many requests we can resume and resume them
+        # Assign released chunks to paused requests.
+        # todo: @shanmugamr, un-pause requests using FIFO, rather than LIFO.
+        if (
+            self.chunk_allocator.chunk_count_avail
+            <= self.paused_request_count + self.gtd_chunk_count
+        ):
+            if active_request_count < self.gtd_request_count:
+                resume_request_count = min(
+                    self.paused_request_count, self.gtd_request_count - active_request_count
+                )
+            else:
+                # If there are more active requests than gtd requests and not enough
+                # chunks available, no requests can be resumed
+                resume_request_count = 0
+        else:
+            # If there are more available chunks than (paused + gtd requests), resume all paused requests
+            resume_request_count = self.paused_request_count
+
+        self.paused_request_count -= resume_request_count
+        active_request_count += resume_request_count
+        assert active_request_count > 0, "active_request_count == %d." % active_request_count
+
+        # 7. We make changes to the request book keeping tesnsors and setup the tokens for next iteration
+        self.total_request_count = active_request_count + self.paused_request_count
+        # All these active requests are in decode phase, so they need only 1 token per request
+        self.active_token_count = active_request_count
+        # Always the first section of token input ids are only used.
+        self.token_to_input_ids[: self.active_token_count] = next_tokens[
+            self.paused_request_count : self.total_request_count
+        ]
+
+        if self.paused_request_count > 0:
+            self.paused_tokens = next_tokens[: self.paused_request_count]
+
+        self.request_kv_length_offsets[self.paused_request_count : self.total_request_count].add_(
+            self.request_query_lengths[self.paused_request_count : self.total_request_count]
+        )
+        self.request_query_lengths[self.paused_request_count : self.total_request_count].fill_(1)
+        self.token_to_pos_ids[: self.active_token_count] = self.request_kv_length_offsets[
+            self.paused_request_count : self.total_request_count
+        ]
+
+        self.request_last_kv_chunk_offset[self.paused_request_count : self.total_request_count] = (
+            self.request_last_kv_chunk_offset[self.paused_request_count : self.total_request_count]
+            + 1
+        ) % self.chunk_size_tokens
+
+        # 8. We resume those requests by assigning chunks and updating bookkeeping tensors
+        if resume_request_count > 0:
+            assert torch.all(
+                self.request_last_kv_chunk_offset[
+                    self.paused_request_count : (self.paused_request_count + resume_request_count)
+                ]
+                == 0
+            ), 'The request_last_kv_chunk_offset should be 0 for the requests that just got resumed this step. '
+
+            chunk_ids = self.chunk_allocator.allocate_memory_chunks(resume_request_count)
+            row_idx = torch.arange(
+                self.paused_request_count,
+                self.paused_request_count + resume_request_count,
+                device=torch.cuda.current_device(),
+            )
+            col_idx = self.request_kv_chunk_counts[
+                self.paused_request_count : (self.paused_request_count + resume_request_count)
+            ]
+            self.request_to_kv_chunk_ids[row_idx, col_idx] = chunk_ids
+            self.request_kv_chunk_counts[
+                self.paused_request_count : (self.paused_request_count + resume_request_count)
+            ] += 1
+            self.request_last_kv_chunk_id[
+                self.paused_request_count : (self.paused_request_count + resume_request_count)
+            ] = chunk_ids
+
+        # 9. We make relevant changes to the token bookkeeping tensors
+        self.token_to_request_idx[: self.active_token_count] = torch.arange(
+            self.paused_request_count, self.total_request_count, device=torch.cuda.current_device()
+        )
+        self.token_to_position_in_request[: self.active_token_count] = (
+            self.request_kv_length_offsets[self.paused_request_count : self.total_request_count]
+        )
+
+        self.token_to_chunk_idx[: self.active_token_count] = self.request_last_kv_chunk_id[
+            self.paused_request_count : self.total_request_count
+        ]
+        self.token_to_local_position_within_kv_chunk[: self.active_token_count] = (
+            self.request_last_kv_chunk_offset[self.paused_request_count : self.total_request_count]
+        )
diff --git a/megatron/core/inference/contexts/static_context.py b/megatron/core/inference/contexts/static_context.py
new file mode 100644
index 000000000..79d865724
--- /dev/null
+++ b/megatron/core/inference/contexts/static_context.py
@@ -0,0 +1,133 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
+    InferenceWrapperConfig,
+)
+
+from .base_context import BaseInferenceContext
+
+
+class StaticInferenceContext(BaseInferenceContext):
+    """Static inference context that is passed to the main model in order
+    to efficiently manage the KV cache during inference.
+
+    Args:
+        max_batch_size (int): Max supported batch size.
+        max_sequence_length (int): Max supported sequence length.
+    """
+
+    def __init__(self, max_batch_size: int, max_sequence_length: int):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        self.max_batch_size = max_batch_size
+        self.current_batch_size = max_batch_size  # Required for bookkeeping variable-sized batches
+        self.sequence_len_offset = 0
+        self.batch_size_offset = 0
+        self.key_value_memory_dict = {}
+        self.decode_mode = False
+        self.materialize_only_last_token_logits = False
+
+    @classmethod
+    def from_config(cls, config: InferenceWrapperConfig) -> "StaticInferenceContext":
+        """Initialize context from a config."""
+        max_batch_size = config.inference_max_requests
+        max_sequence_length = config.inference_max_seq_length
+        return cls(max_batch_size, max_sequence_length)
+
+    def swap_key_value_dict(self, batch_idx):
+        "swap between batches"
+        if len(self.key_value_memory_dict) == 0:
+            raise ValueError("should not swap when dict in empty")
+
+        for layer_number in self.key_value_memory_dict.keys():
+            inference_key_memory, inference_value_memory = self.key_value_memory_dict[layer_number]
+            assert (
+                len(batch_idx) == inference_key_memory.shape[1]
+            )  # make sure batch size is the same
+            new_inference_key_memory = inference_key_memory[:, batch_idx]
+            new_inference_value_memory = inference_value_memory[:, batch_idx]
+            self.key_value_memory_dict[layer_number] = (
+                new_inference_key_memory,
+                new_inference_value_memory,
+            )
+
+    def enable_prefill_mode(self):
+        """
+        Indicates the generation loop is in the prefill phase (still processing
+        input prompt tokens). This should be enabled if the generation loop is
+        encoding prompt tokens for *any* request in a batch.
+        """
+        self.decode_mode = False
+
+    def enable_decode_mode(self):
+        """
+        Indicates the generation loop is in the decode phase (generating new output
+        tokens). This should only be enabled if the generation loop has fully encoded
+        the prompts for *all* requests in a batch.
+        """
+        self.decode_mode = True
+
+    def is_decode_only(self):
+        """Functional access to `.decode_mode`, to match dynamic context."""
+        return self.decode_mode
+
+    def reset(self):
+        """Resets the inference state for a new batch."""
+        self.current_batch_size = self.max_batch_size
+        self.sequence_len_offset = 0
+        self.batch_size_offset = 0
+        self.enable_prefill_mode()
+
+    def __str__(self):
+        return (
+            f"StaticInferenceContext(max_seq_len = {self.max_sequence_length}, "
+            f"max_batch_size = {self.max_batch_size}, "
+            f"current_batch_size = {self.current_batch_size}, "
+            f"sequence_len_offset = {self.sequence_len_offset}, "
+            f"batch_size_offset = {self.batch_size_offset}, "
+            f"key_value_memory_dict = {self.key_value_memory_dict.keys()})"
+            f"decode_mode = {self.decode_mode}"
+            f"materialize_only_last_token_logits = {self.materialize_only_last_token_logits}"
+        )
+
+    def __eq__(self, other):
+
+        if id(self) == id(other):
+            return True
+
+        if not isinstance(other, StaticInferenceContext):
+            return False
+
+        # Check all attributes match
+        basic_attrs = [
+            'max_sequence_length',
+            'max_batch_size',
+            'current_batch_size',
+            'sequence_len_offset',
+            'batch_size_offset',
+            'decode_mode',
+            'materialize_only_last_token_logits',
+        ]
+
+        if not all(hasattr(other, attr) for attr in basic_attrs):
+            return False
+
+        # Check dictionary keys match; i.e. the same number of layers are cached
+        if self.key_value_memory_dict.keys() != other.key_value_memory_dict.keys():
+            return False
+
+        # Check each tensor tuple in the dictionary
+        for key in self.key_value_memory_dict:
+            self_tensors = self.key_value_memory_dict[key]
+            other_tensors = other.key_value_memory_dict[key]
+
+            # Compare each key, value tensor in the tuple
+            for self_tensor, other_tensor in zip(self_tensors, other_tensors):
+                if (
+                    self_tensor.data_ptr() != other_tensor.data_ptr()
+                    or self_tensor.shape != other_tensor.shape
+                ):
+                    return False
+
+    def is_static_batching(self):
+        return True
diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index cb26be122..898e7f0b5 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -13,7 +13,10 @@ from megatron.core.transformer.enums import AttnBackend
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import make_tp_sharded_tensor_for_checkpoint
-
+try:
+    from megatron.core.extensions.transformer_engine import te_parallel_cross_entropy
+except:
+    te_parallel_cross_entropy = None
 
 class LanguageModule(MegatronModule):
     """Base language module that has common helper functions used across GPT, BERT etc.
@@ -76,7 +79,10 @@ class LanguageModule(MegatronModule):
         # [b s] => [s b]
         labels = labels.transpose(0, 1).contiguous()
         if self.config.cross_entropy_loss_fusion:
-            loss = fused_vocab_parallel_cross_entropy(logits, labels)
+            # loss = fused_vocab_parallel_cross_entropy(logits, labels)
+            # HACK(yehua.zhang): replace fuse cross entropy to triton
+            labels = torch.as_strided(labels, labels.size(), (labels.size()[1], 1))
+            loss = te_parallel_cross_entropy(logits, labels)
         else:
             loss = tensor_parallel.vocab_parallel_cross_entropy(logits, labels)
 
diff --git a/megatron/core/optimizer/clip_grads.py b/megatron/core/optimizer/clip_grads.py
index 0f33f919b..f5cb2c49d 100644
--- a/megatron/core/optimizer/clip_grads.py
+++ b/megatron/core/optimizer/clip_grads.py
@@ -87,7 +87,7 @@ def get_grad_norm_fp32(
     # Calculate norm.
     if norm_type == inf:
         total_norm = max(grad.abs().max() for grad in grads_for_norm)
-        total_norm_cuda = torch.tensor([float(total_norm)], dtype=torch.float, device='cuda')
+        total_norm_cuda = torch.tensor([float(total_norm)], dtype=torch.float, device='musa')
         # Take max across all data-parallel GPUs if using FSDP and then all model-parallel GPUs.
         if data_parallel_group:
             torch.distributed.all_reduce(
@@ -100,7 +100,7 @@ def get_grad_norm_fp32(
 
     else:
         if norm_type == 2.0:
-            dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device='cuda')
+            dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device='musa')
             # Use apex's multi-tensor applier for efficiency reasons.
             # Multi-tensor applier takes a function and a list of list
             # and performs the operation on that list all in one kernel.
@@ -112,7 +112,7 @@ def get_grad_norm_fp32(
                     False,  # no per-parameter norm
                 )
             else:
-                grad_norm = torch.tensor([0], dtype=torch.float, device='cuda')
+                grad_norm = torch.tensor([0], dtype=torch.float, device='musa')
             # Since we will be summing across data parallel groups,
             # we need the pow(norm-type).
             total_norm = grad_norm**norm_type
@@ -171,7 +171,7 @@ def clip_grad_by_total_norm_fp32(
     # Scale.
     clip_coeff = max_norm / (total_norm + 1.0e-6)
     if clip_coeff < 1.0:
-        dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device='cuda')
+        dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device='musa')
         multi_tensor_applier(
             multi_tensor_scale_impl, dummy_overflow_buf, [grads, grads], clip_coeff
         )
@@ -203,7 +203,7 @@ def count_zeros_fp32(
     #   - grad should not be none
     #   - parameter should not be shared
     #   - should not be a replica due to tensor model parallelism
-    total_num_zeros = torch.tensor([0.0], dtype=torch.float, device='cuda')
+    total_num_zeros = torch.tensor([0.0], dtype=torch.float, device='musa')
     data_parallel_group = None
     for param in parameters:
         grad_attr = "decoupled_grad" if use_decoupled_grad else "grad"
diff --git a/megatron/core/process_groups_config.py b/megatron/core/process_groups_config.py
new file mode 100644
index 000000000..da39a27c3
--- /dev/null
+++ b/megatron/core/process_groups_config.py
@@ -0,0 +1,178 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+"""Dataclasses for organizing model parallelism and gradient communication process groups."""
+
+from dataclasses import dataclass, field, fields
+from typing import List, Optional
+
+import torch
+
+from megatron.core import parallel_state
+
+
+@dataclass
+class ModelCommProcessGroups:
+    """Process groups for transformer model parallelism.
+
+    Fields use init=False and must be set after instance creation.
+
+    Args:
+        tp: Tensor parallel process group
+        pp: Pipeline parallel process group
+        mp: Model parallel group (tensor + pipeline)
+        embd: Embedding process group
+        pos_embd: Position embedding process group
+        cp: Context parallel process group
+        tp_cp: Tensor and context parallel group
+        hcp: Hierarchical context parallel groups
+        ep: Expert model parallel group
+        expt_tp: Expert tensor parallel group
+        tp_ep: Tensor and expert parallel group
+        tp_ep_pp: Tensor, expert, and pipeline parallel group
+        expt_dp: Expert data parallel group
+    Example:
+        # Create instance and set needed process groups
+        model_pgs = ModelCommProcessGroups()
+        model_pgs.tp = tp_group
+        model_pgs.pp = pp_group
+
+        # Pass to model components
+        model = TransformerModel(..., process_groups=model_pgs)
+    """
+
+    # _TENSOR_MODEL_PARALLEL_GROUP
+    tp: torch.distributed.ProcessGroup = field(init=False)
+
+    # _PIPELINE_MODEL_PARALLEL_GROUP
+    pp: torch.distributed.ProcessGroup = field(init=False)
+
+    # _MODEL_PARALLEL_GROUP
+    mp: torch.distributed.ProcessGroup = field(init=False)
+
+    # _EMBEDDING_GROUP
+    embd: torch.distributed.ProcessGroup = field(init=False)
+
+    # _POSITION_EMBEDDING_GROUP
+    pos_embd: torch.distributed.ProcessGroup = field(init=False)
+
+    # _CONTEXT_PARALLEL_GROUP
+    cp: torch.distributed.ProcessGroup = field(init=False)
+
+    # _TENSOR_AND_CONTEXT_PARALLEL_GROUP
+    tp_cp: torch.distributed.ProcessGroup = field(init=False)
+
+    # _HIERARCHICAL_CONTEXT_PARALLEL_GROUPS
+    hcp: List[torch.distributed.ProcessGroup] = field(init=False)
+
+    # _EXPERT_MODEL_PARALLEL_GROUP
+    ep: torch.distributed.ProcessGroup = field(init=False)
+
+    # _EXPERT_TENSOR_PARALLEL_GROUP
+    expt_tp: torch.distributed.ProcessGroup = field(init=False)
+
+    # _EXPERT_TENSOR_AND_MODEL_PARALLEL_GROUP
+    tp_ep: torch.distributed.ProcessGroup = field(init=False)
+
+    # _EXPERT_TENSOR_MODEL_PIPELINE_PARALLEL_GROUP
+    tp_ep_pp: torch.distributed.ProcessGroup = field(init=False)
+
+    # MoE layers need expt_dp group for sharded state dict
+    # we need this workaround until distributed checkpoint is refactored
+    # to have sharded_state_dict can take the PG and pass it down
+    # TODO (Hepteract): remove this once distributed checkpoint is refactored
+    # _EXPERT_DATA_PARALLEL_GROUP
+    expt_dp: torch.distributed.ProcessGroup = field(init=False)
+
+    def __init__(self, **kwargs):
+        for key in kwargs:
+            if key in [field.name for field in fields(self)]:
+                setattr(self, key, kwargs[key])
+            else:
+                raise ValueError(f"Unknown attribute: {key}")
+
+    @classmethod
+    def use_mpu_process_groups(cls, required_pgs: Optional[List[str]] = None):
+        """
+        Use the default process groups from parallel_state.
+
+        Args:
+            required_pgs (List[str], optional): List of process group names to initialize.
+                If None, pull all default process groups. Each string should correspond to
+                one of the dataclass process group attributes.
+        """
+        # Get all available process groups
+        all_pgs = {field.name for field in fields(cls)}
+
+        # If no specific process groups requested, use all
+        if required_pgs is None:
+            required_pgs = list(all_pgs)
+
+        # Validate requested process groups
+        invalid_pgs = [pg for pg in required_pgs if pg not in all_pgs]
+        if invalid_pgs:
+            raise ValueError(f"Invalid process groups requested: {invalid_pgs}")
+
+        # Mapping of attribute names to their initialization functions
+        pg_to_func = {
+            'tp': parallel_state.get_tensor_model_parallel_group,
+            'pp': parallel_state.get_pipeline_model_parallel_group,
+            'mp': parallel_state.get_model_parallel_group,
+            'cp': parallel_state.get_context_parallel_group,
+            'tp_cp': parallel_state.get_tensor_and_context_parallel_group,
+            'hcp': parallel_state.get_hierarchical_context_parallel_groups,
+            'ep': parallel_state.get_expert_model_parallel_group,
+            'expt_tp': parallel_state.get_expert_tensor_parallel_group,
+            'tp_ep': parallel_state.get_expert_tensor_and_model_parallel_group,
+            'tp_ep_pp': parallel_state.get_expert_tensor_model_pipeline_parallel_group,
+            'embd': parallel_state.get_embedding_group,
+            'pos_embd': parallel_state.get_position_embedding_group,
+            # TODO (Hepteract): remove this once distributed checkpoint is refactored
+            'expt_dp': parallel_state.get_expert_data_parallel_group,
+        }
+
+        # Build initialization dict by calling appropriate parallel_state get_foo_group
+        init_dict = {pg: pg_to_func[pg](False) for pg in required_pgs}
+
+        return cls(**init_dict)
+
+
+@dataclass
+class GradCommProcessGroups:
+    """Process groups for gradient communication in distributed training.
+
+    Fields use init=False and must be set after instance creation.
+
+    Args:
+        dp: Data parallel process group
+        dp_cp: Data and context parallel group
+        expt_dp: Expert data parallel group
+        intra_dp_cp: Intra partial data parallel group
+        intra_expt_dp: Intra partial expert data parallel group
+        inter_dist_opt: Inter distributed optimizer instance group
+
+    Example:
+        # Create instance and set needed process groups
+        grad_pgs = GradCommProcessGroups()
+        grad_pgs.dp = dp_group
+
+        # Pass to distributed data parallel wrapper
+        ddp_model = DistributedDataParallel(..., process_groups=grad_pgs)
+    """
+
+    # _DATA_PARALLEL_GROUP
+    dp: torch.distributed.ProcessGroup = field(init=False)
+
+    # _DATA_PARALLEL_GROUP_WITH_CP
+    dp_cp: torch.distributed.ProcessGroup = field(init=False)
+
+    # _EXPERT_DATA_PARALLEL_GROUP
+    expt_dp: torch.distributed.ProcessGroup = field(init=False)
+
+    # _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP
+    intra_dp_cp: torch.distributed.ProcessGroup = field(init=False)
+
+    # _INTRA_EXPERT_DATA_PARALLEL_GROUP
+    intra_expt_dp: torch.distributed.ProcessGroup = field(init=False)
+
+    # _INTER_DISTRIBUTED_OPTIMIZER_INSTANCE_GROUP
+    inter_dist_opt: torch.distributed.ProcessGroup = field(init=False)
diff --git a/megatron/core/tensor_parallel/cross_entropy.py b/megatron/core/tensor_parallel/cross_entropy.py
index 27c8f0634..bd5890838 100644
--- a/megatron/core/tensor_parallel/cross_entropy.py
+++ b/megatron/core/tensor_parallel/cross_entropy.py
@@ -59,6 +59,7 @@ class VocabParallelCrossEntropy:
         predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
         predicted_logits_1d = predicted_logits_1d.clone().contiguous()
         predicted_logits = predicted_logits_1d.view_as(target)
+        # print(f'target_mask is {target_mask.shape}, {target_mask.dtype}, predicted_logits is {predicted_logits.shape}, {predicted_logits.dtype}')
         predicted_logits[target_mask] = 0.0
 
         exp_logits = vocab_parallel_logits
diff --git a/megatron/core/tensor_parallel/mappings.py b/megatron/core/tensor_parallel/mappings.py
index cdd720687..66bf20839 100644
--- a/megatron/core/tensor_parallel/mappings.py
+++ b/megatron/core/tensor_parallel/mappings.py
@@ -439,6 +439,10 @@ class _AllToAll(torch.autograd.Function):
                 dtype=input.dtype,
                 device=torch.cuda.current_device(),
             )
+        # print(f'output_split_sizes is {output_split_sizes}')
+        # print(f'input_split_sizes is {input_split_sizes}')
+        # print(f'output is {output.shape}')
+        # print(f'input is {input.shape}')
         torch.distributed.all_to_all_single(
             output,
             input,
diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py
index fc0cf3e62..2a8579911 100644
--- a/megatron/core/transformer/moe/experts.py
+++ b/megatron/core/transformer/moe/experts.py
@@ -682,11 +682,12 @@ class TEGroupedMLP(MegatronModule):
             output (torch.Tensor): The output of the local experts.
         """
         tokens_per_expert = tokens_per_expert.tolist()
-        if self.config.fp8:
-            actual_tokens_per_expert = tokens_per_expert
-            permuted_local_hidden_states, tokens_per_expert = self.fp8_padding(
-                permuted_local_hidden_states, tokens_per_expert
-            )
+        # TODO(yehua.zhang): musa groupgemm do not need to padding
+        # if self.config.fp8:
+        #     actual_tokens_per_expert = tokens_per_expert
+        #     permuted_local_hidden_states, tokens_per_expert = self.fp8_padding(
+        #         permuted_local_hidden_states, tokens_per_expert
+        #     )
 
         intermediate_parallel, bias_parallel = self.linear_fc1(
             permuted_local_hidden_states, tokens_per_expert
@@ -734,8 +735,9 @@ class TEGroupedMLP(MegatronModule):
         output, output_bias = self.linear_fc2(intermediate_parallel, tokens_per_expert)
 
         # upad and concat the output
-        if self.config.fp8:
-            output = self.fp8_unpadding(output, actual_tokens_per_expert)
+        # TODO(yehua.zhang): musa groupgemm do not need to unpadding
+        # if self.config.fp8:
+        #     output = self.fp8_unpadding(output, actual_tokens_per_expert)
 
         return output, output_bias
 
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index cf7cf2b43..87886e386 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -354,7 +354,8 @@ def sort_chunks_by_idxs(
     input: torch.Tensor, split_sizes: torch.Tensor, sorted_idxs: torch.Tensor, fused: bool = False
 ):
     """Split and sort the input tensor based on the split_sizes and sorted indices."""
-    if fused:
+    # TODO(yehua.zhang) optimize the sort chunk kernel
+    if False:#fused:
         if not HAVE_TE or fused_sort_chunks_by_index is None:
             raise ValueError(
                 "fused_sort_chunks_by_index is not available. Please install TE >= 2.1.0."
@@ -403,7 +404,8 @@ def group_limited_topk(
         Tuple[torch.Tensor, torch.Tensor]: Probs and indices tensor.
     """
     # Organize the experts into groups
-    group_scores = scores.view(num_tokens, num_groups, -1).topk(2, dim=-1)[0].sum(dim=-1)
+    # TODO(yehua.zhang) delete the max
+    group_scores = scores.view(num_tokens, num_groups, -1).max(dim=-1).values #.topk(2, dim=-1)[0].sum(dim=-1)
     group_idx = torch.topk(group_scores, k=group_topk, dim=-1, sorted=False)[1]
     group_mask = torch.zeros_like(group_scores)
     group_mask.scatter_(1, group_idx, 1)
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 5965c16dc..75023684a 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -112,9 +112,11 @@ class TopKRouter(Router):
                 torch.zeros(self.config.num_moe_experts, dtype=torch.float32),
                 persistent=False,
             )
+            self.local_tokens_per_expert = self.local_tokens_per_expert.cuda()
             self.register_buffer(
                 'expert_bias', torch.zeros(self.config.num_moe_experts, dtype=torch.float32)
             )
+            self.expert_bias = self.expert_bias.cuda()
         else:
             self.local_tokens_per_expert = None
             self.expert_bias = None
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index d92d86e79..36bbba732 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -755,8 +755,8 @@ class TransformerConfig(ModelParallelConfig):
                     "apply_rope_fusion is not available. Please install TE >= 1.4 or Apex."
                 )
 
-            if self.multi_latent_attention:
-                raise ValueError("multi_latent_attention does not support apply_rope_fusion.")
+            # if self.multi_latent_attention:
+            #     raise ValueError("multi_latent_attention does not support apply_rope_fusion.")
 
         if self.multi_latent_attention and self.rotary_interleaved:
             raise ValueError("rotary_interleaved does not work with multi_latent_attention.")
@@ -857,12 +857,12 @@ class TransformerConfig(ModelParallelConfig):
                 fused_unpermute,
             )
 
-            if (
-                fused_permute is None
-                or fused_sort_chunks_by_index is None
-                or fused_unpermute is None
-            ):
-                raise ValueError("fused permutation is not available. Please install TE >= 2.1.0.")
+            # if (
+            #     fused_permute is None
+            #     or fused_sort_chunks_by_index is None
+            #     or fused_unpermute is None
+            # ):
+            #     raise ValueError("fused permutation is not available. Please install TE >= 2.1.0.")
 
         if self.cp_comm_type is not None:
             if isinstance(self.cp_comm_type, list):
diff --git a/megatron/core/utils.py b/megatron/core/utils.py
index f80e3e04d..e89aed190 100644
--- a/megatron/core/utils.py
+++ b/megatron/core/utils.py
@@ -62,6 +62,24 @@ def get_torch_version():
     return _torch_version
 
 
+def deprecate_inference_params(inference_context, inference_params):
+    """Print warning for deprecated `inference_params`."""
+    if inference_context is None and inference_params is not None:
+        warnings.warn(
+            "`inference_params` renamed to `inference_context`, and will be "
+            "removed in `megatron-core` 0.13."
+        )
+        return inference_params
+    return inference_context
+
+
+def is_fa_min_version(version, check_equality=True):
+    """Check if minimum version of `flash-attn` is installed."""
+    if check_equality:
+        return get_fa_version() >= PkgVersion(version)
+    return get_fa_version() > PkgVersion(version)
+
+
 def get_te_version():
     """Get TE version from __version__; if not available use pip's. Use caching."""
 
@@ -688,6 +706,26 @@ def local_multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale):
         dst.copy_(src * scale)
 
 
+class WrappedTensor:
+    """
+    A wrapper for tensors that enables caller functions to pass an indirect reference
+    to callee functions. By wrapping the tensor, the caller's direct reference is removed,
+    allowing the tensor to be garbage collected once the callee unwraps and frees it.
+    """
+
+    def __init__(self, tensor: torch.Tensor):
+        self._wrapper = [tensor]
+
+    def unwrap(self):
+        """
+        Returns the wrapped tensor while deleting the internal reference.
+        Can only be called once.
+        """
+        if len(self._wrapper) == 0:
+            raise RuntimeError(f"WrappedTensor has already been unwrapped")
+        return self._wrapper.pop(0)
+
+
 class _ValueWithRank:
     """This is an internal class, not for use outside this module
 
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 06f06dae0..8d2dac15f 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1083,8 +1083,10 @@ def _add_network_size_args(parser):
                        help='Maximum number of position embeddings to use. '
                        'This is the size of position embedding.')
     group.add_argument('--position-embedding-type', type=str, default='learned_absolute',
-                        choices=['learned_absolute', 'rope', 'relative', 'none'],
+                        choices=['learned_absolute', 'rope', 'relative', 'none', 'mrope'],
                         help='Position embedding type.')
+    group.add_argument('--mrope-section', nargs='+', type=int, default=None,
+                       help='Multimodal rope section is for channel dimension, empty by default.')
     group.add_argument('--relative-attention-num-buckets', type=int, default=32,
                         help='Number of buckets for relative position embeddings.')
     group.add_argument('--relative-attention-max-distance', type=int, default=128,
@@ -2012,6 +2014,8 @@ def _add_tokenizer_args(parser):
     group.add_argument('--tokenizer-type', type=str,
                        default=None,
                        choices=['BertWordPieceLowerCase',
+                                'Qwen2TokenizerFS',
+                                'Qwen2VLTokenizer',
                                 'BertWordPieceCase',
                                 'GPT2BPETokenizer',
                                 'SentencePieceTokenizer',
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index 749563caf..9fe2218ea 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -1232,6 +1232,7 @@ def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', stri
     set_checkpoint_version(state_dict.get('checkpoint_version', 0))
 
     # Set iteration.
+    #print(f'info+ {args.finetune} {list(state_dict.keys())}')
     if args.finetune or release:
         iteration = 0
     else:
@@ -1267,7 +1268,8 @@ def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', stri
     strict = False if args.retro_add_retriever else strict
     if not skip_load_to_model_and_opt:
         if len(ddp_model) == 1:
-            ddp_model[0].load_state_dict(state_dict['model'], strict=strict)
+
+            ddp_model[0].load_state_dict(state_dict['model'], strict=False)
         else:
             for i in range(len(ddp_model)):
                 mpu.set_virtual_pipeline_model_parallel_rank(i)
@@ -1283,6 +1285,7 @@ def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', stri
         try:
             # Load state dict.
             if not skip_load_to_model_and_opt and optimizer is not None and not optimizer.is_stub_optimizer:
+                # print(state_dict.keys())
                 optimizer.load_state_dict(state_dict['optimizer'])
 
             # Load distributed optimizer's custom parameter state.
@@ -1380,7 +1383,6 @@ def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', stri
         # Notify FT that a checkpoint was loaded.
         is_local_chkpt = (ckpt_type == CheckpointType.LOCAL)
         ft_integration.on_checkpoint_loaded(is_local_chkpt=is_local_chkpt)
-
     return iteration, num_floating_point_operations_so_far
 
 
diff --git a/megatron/training/tokenizer/tokenizer.py b/megatron/training/tokenizer/tokenizer.py
index 620a0cbbb..2e35ed87a 100644
--- a/megatron/training/tokenizer/tokenizer.py
+++ b/megatron/training/tokenizer/tokenizer.py
@@ -65,6 +65,14 @@ def build_tokenizer(args, **kwargs):
     elif args.tokenizer_type == 'NullTokenizer':
         assert args.vocab_size is not None
         tokenizer = _NullTokenizer(args.vocab_size)
+    elif args.tokenizer_type == "Qwen2TokenizerFS":
+        assert args.tokenizer_model is not None
+        tokenizer = _Qwen2TokenizerFS(args.tokenizer_model, args)
+    elif args.tokenizer_type == 'Qwen2VLTokenizer':
+        assert args.tokenizer_model is not None
+        tokenizer = _Qwen2VLTokenizer(args.tokenizer_model, args.extra_vocab_size)
+        args.padded_vocab_size = tokenizer.vocab_size # no padding
+
     elif args.tokenizer_type == "MultimodalTokenizer":
         try:
             import transformers
@@ -129,7 +137,7 @@ class _HuggingFaceTokenizer(MegatronTokenizer):
             )
 
         # TODO(bnorick): download tokenizer once to lustre and use force offline to make sure all tasks read it from there
-        self._tokenizer = transformers.AutoTokenizer.from_pretrained(
+        self._tokenizer = transformers.AutoTokenizer.from_pretrained(trust_remote_code=True,
             pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs
         )
         self._vocab = self._tokenizer.get_vocab()
@@ -835,3 +843,160 @@ class _NullTokenizer(MegatronTokenizer):
     @property
     def additional_special_tokens_ids(self):
         return None
+
+
+class _HFTokenizerFS(MegatronTokenizer):
+    """Huggingface tokenizer."""
+
+    def __init__(self, tokenizer_path):
+        name = 'HFTokenizer'
+        super().__init__(name)
+
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
+
+        self.eod_id = self.tokenizer.eos_token_id
+        self.cls_id = self.tokenizer.bos_token_id
+        self.pad_id = self.tokenizer.pad_token_id
+
+        self._inv_vocab = None
+
+    @property
+    def vocab_size(self):
+        return self.tokenizer.vocab_size
+
+    @property
+    def vocab(self):
+        return self.tokenizer.get_vocab()
+
+    @property
+    def inv_vocab(self):
+        vocab = self.vocab()
+        if self._inv_vocab is None:
+            self._inv_vocab = {v: k for k, v in vocab.items()}
+        return self._inv_vocab
+
+    def tokenize(self, text):
+        return self.tokenizer.encode(text)
+
+    def detokenize(self, token_ids):
+        return self.tokenizer.decode(token_ids)
+
+    @property
+    def eod(self):
+        return self.eod_id
+
+    @property
+    def cls(self):
+        return self.cls_id
+
+    @property
+    def pad(self):
+        return self.pad_id
+
+
+
+class _Qwen2VLTokenizer(MegatronTokenizer):
+    def __init__(self, tokenizer_path, extra_vocab_size):
+        super().__init__(tokenizer_path)
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_path,
+            padding_side="right",
+            use_fast=False,
+            trust_remote_code=True
+        )
+        self.extra_vocab_size = extra_vocab_size
+        self.special_tokens_map = {k:v for k, v in zip(self.tokenizer.all_special_tokens, self.tokenizer.all_special_ids)}
+        self.image_token = '<|image_pad|>'
+        self.video_token = '<|video_pad|>'
+        self.vision_start_token = '<|vision_start|>'
+        self.vision_end_token = '<|vision_end|>'
+
+        from transformers import AutoProcessor
+        proc = AutoProcessor.from_pretrained(
+            tokenizer_path,
+            use_fast=False,
+            trust_remote_code=True
+        )
+        # NOTE: In Qwen2-VL, template in chat_template.json is same within tokenizer_config.json and both can be used.
+        # However, in Qwen 2.5-VL, the two templates are different and only the one in chat_template.json is OK.
+        self.chat_template = proc.chat_template
+
+    def __call__(self, text, return_tensors=None,
+                    padding=None, max_length=None, truncation=None, add_special_tokens=None):
+
+        return self.tokenizer(text, return_tensors=return_tensors, padding=padding,
+                max_length=max_length, truncation=truncation, add_special_tokens=add_special_tokens)
+
+    def apply_chat_template(self, conversations, tokenize:bool=True, **kwargs):
+        return self.tokenizer.apply_chat_template(conversations, tokenize=tokenize, chat_template=self.chat_template, **kwargs)
+
+    @property
+    def vocab_size(self):
+        return len(self.tokenizer.encoder) + self.extra_vocab_size
+
+    @property
+    def vocab(self):
+        return self.tokenizer.encoder
+
+    @property
+    def inv_vocab(self):
+        return self.tokenizer.decoder
+
+    def tokenize(self, text):
+        return self.tokenizer.encode(text)
+
+    def detokenize(self, token_ids):
+        return self.tokenizer.decode(token_ids)
+
+    @property
+    def eod(self):
+        return self.tokenizer.eos_token_id
+
+    @property
+    def eos_token(self):
+        return self.tokenizer.eos_token
+
+    @property
+    def pad_token_id(self):
+        return self.tokenizer.pad_token_id
+
+    @property
+    def eos_token_id(self):
+        return self.tokenizer.eos_token_id
+
+    @property
+    def image_token_id(self):
+        return self.special_tokens_map[self.image_token]
+
+    @property
+    def video_token_id(self):
+        return self.special_tokens_map[self.video_token]
+
+    @property
+    def vision_start_token_id(self):
+        return self.special_tokens_map[self.vision_start_token]
+
+    @property
+    def vision_end_token_id(self):
+        return self.special_tokens_map[self.vision_end_token]
+
+    def encode(self, x):
+        return self.tokenizer.encode(x)
+
+
+class _Qwen2TokenizerFS(_HFTokenizerFS):
+    """Adapted Qwen tokenizer."""
+
+    def __init__(self, tokenizer_path, args):
+        super().__init__(tokenizer_path)
+        self.eod_id = self.tokenizer.encode('<|extra_204|>')[0]
+        self.cls_id = self.tokenizer.encode('<|extra_203|>')[0]
+        self.pad_id = self.tokenizer.encode('<|endoftext|>')[0]
+        assert args.vocab_size is not None
+        self._vocab_size = args.vocab_size
+
+    @property
+    def vocab_size(self):
+        return self._vocab_size
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 7cf6fcbd0..768dcd2c2 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -52,6 +52,7 @@ from megatron.core.rerun_state_machine import (
     RerunDataIterator,
     RerunMode,
 )
+from megatron.core.fp8_utils import get_fp8_scale_and_amax
 from megatron.training.initialize import initialize_megatron
 from megatron.training.initialize import write_args_to_tensorboard
 from megatron.training.initialize import set_jit_fusion_options
@@ -582,9 +583,9 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap
     # GPU allocation.
     # For FSDP2, we don't allocate GPU memory here. We allocate GPU memory
     # in the fully_shard function of FSDP2 instead.
-    if not (args.use_torch_fsdp2 and args.use_cpu_initialization) and not args.init_model_with_meta_device:
-        for model_module in model:
-            model_module.cuda(torch.cuda.current_device())
+    # if not (args.use_torch_fsdp2 and args.use_cpu_initialization) and not args.init_model_with_meta_device:
+    #     for model_module in model:
+    #         model_module.cuda(torch.cuda.current_device())
 
     # Fp16 conversion.
     if args.fp16 or args.bf16:
@@ -595,15 +596,19 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap
     # param) to its amax_history. The following logic will correct the amax_history back.
     for model_module in model:
         for param in model_module.parameters():
-            if is_float8tensor(param) and param._fp8_meta is not None:
-                fp8_meta = param._fp8_meta['scaling_fwd']
-                fp8_meta_index = param._fp8_meta_index
-                if hasattr(param, 'get_high_precision_init_val'):
-                    fp8_meta.amax_history[0][fp8_meta_index].copy_(
-                        param.get_high_precision_init_val().abs().max()
-                    )
-                else:
-                    fp8_meta.amax_history[0][fp8_meta_index] = 0
+            if is_float8tensor(param): #and param._fp8_meta is not None:
+                # TODO(yehua.zhang): move the fp8_meta_index
+                scale, amax = get_fp8_scale_and_amax(param)
+                # amax = 0
+                amax[0].zero_()
+                # fp8_meta = param._fp8_meta['scaling_fwd']
+                # fp8_meta_index = param._fp8_meta_index
+                # if hasattr(param, 'get_high_precision_init_val'):
+                #     fp8_meta.amax_history[0][fp8_meta_index].copy_(
+                #         param.get_high_precision_init_val().abs().max()
+                #     )
+                # else:
+                #     fp8_meta.amax_history[0][fp8_meta_index] = 0
 
     if wrap_with_ddp:
         if args.use_torch_fsdp2:
