diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index fdc6bc24..09c73151 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -15,7 +15,7 @@
 workflow:
   rules:
     # Do not trigger for forks
-    - if: $CI_PROJECT_NAMESPACE != "ADLR" || $CI_MERGE_REQUEST_PROJECT_PATH != "ADLR/megatron-lm"
+    - if: $CI_PROJECT_NAMESPACE != "ADLR" || ($CI_PIPELINE_SOURCE == "merge_request_event" && $CI_MERGE_REQUEST_PROJECT_PATH != "ADLR/megatron-lm")
       when: never
 
     # ci-branches only for schedule
@@ -69,7 +69,9 @@ workflow:
     # For MRs with nightly
     - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result' && $CI_MERGE_REQUEST_LABELS =~ /Run nightly/
       variables:
-        UNIT_TEST: "no"
+        UNIT_TEST: "yes"
+        UNIT_TEST_REPEAT: 1
+        UNIT_TEST_TIMEOUT: 15
         INTEGRATION_TEST: "no"
         FUNCTIONAL_TEST: "yes"
         FUNCTIONAL_TEST_SCOPE: nightly
@@ -83,7 +85,9 @@ workflow:
     # For MRs with weekly
     - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result' && $CI_MERGE_REQUEST_LABELS =~ /Run weekly/
       variables:
-        UNIT_TEST: "no"
+        UNIT_TEST: "yes"
+        UNIT_TEST_REPEAT: 1
+        UNIT_TEST_TIMEOUT: 15
         INTEGRATION_TEST: "no"
         FUNCTIONAL_TEST: "yes"
         FUNCTIONAL_TEST_SCOPE: weekly
diff --git a/.gitlab/stages/00.pre.yml b/.gitlab/stages/00.pre.yml
index cf44ffe0..c4e92a35 100644
--- a/.gitlab/stages/00.pre.yml
+++ b/.gitlab/stages/00.pre.yml
@@ -59,6 +59,7 @@ pre:create_ci_branches:
       - branch: ci-nightly
       - branch: ci-weekly
       - branch: ci-pre-release
+      - branch: ci-review-reminder
   tags:
     - arch/amd64
     - env/prod
@@ -214,8 +215,9 @@ pre:check_milestone:
       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
+        LATEST_MILESTONE=$(curl --header "PRIVATE-TOKEN: ${PROJECT_ACCESS_TOKEN_MCORE}" --url "https://${GITLAB_ENDPOINT}/api/v4/projects/${CI_PROJECT_ID}/milestones?state=active&order_by=due_date&sort=desc" | jq '.[0].id')
+        curl --request PUT --header "PRIVATE-TOKEN: ${PROJECT_ACCESS_TOKEN_MCORE}" --url "https://${GITLAB_ENDPOINT}/api/v4/projects/${CI_PROJECT_ID}/merge_requests/${CI_MERGE_REQUEST_IID}" --data "milestone_id=${LATEST_MILESTONE}"
+        echo "Applied latest milestone (ID: ${LATEST_MILESTONE}) to this MR"
       fi
 
 pre:check_status_of_main:
diff --git a/.gitlab/stages/01.build.yml b/.gitlab/stages/01.build.yml
index 81fb6ddd..db76847c 100644
--- a/.gitlab/stages/01.build.yml
+++ b/.gitlab/stages/01.build.yml
@@ -1,5 +1,10 @@
+.build_rules:
+  rules:
+    - when: on_success
+  stage: test
+
 .build_image:
-  extends: [.test_rules, .dind_rules]
+  extends: [.build_rules, .dind_rules]
   stage: build
   tags:
     - arch/amd64
@@ -53,8 +58,6 @@ test:build_image:
       - IMAGE: UTILITY_IMAGE
         FILE: Dockerfile.linting
         BASE_IMAGE: python:3.10
-  rules:
-    - when: always
 
 test:build_nemo_image:
   extends: [.build_image]
@@ -62,3 +65,6 @@ test:build_nemo_image:
     IMAGE: CI_NEMO_IMAGE
     FILE: Dockerfile.ci.dev
     BASE_IMAGE: nvcr.io/nvidian/nemo:nightly
+  rules:
+    - if: $FUNCTIONAL_TEST == "yes" || $INTEGRATION_TEST == "yes"
+      when: on_success
diff --git a/.gitlab/stages/02.test.yml b/.gitlab/stages/02.test.yml
index d0c51c08..bde211d9 100644
--- a/.gitlab/stages/02.test.yml
+++ b/.gitlab/stages/02.test.yml
@@ -1,6 +1,6 @@
 .test_rules:
   rules:
-    - if: $PUBLISH == "yes" && $PUBLISH_SCOPE == "review-reminder"
+    - if: $PUBLISH == "yes"
       when: never
     - when: on_success
   stage: test
diff --git a/.gitlab/stages/03.integration-tests.yml b/.gitlab/stages/03.integration-tests.yml
index 109a7180..fcdd57f2 100644
--- a/.gitlab/stages/03.integration-tests.yml
+++ b/.gitlab/stages/03.integration-tests.yml
@@ -15,8 +15,32 @@ include:
     ref: main
     file: downstreams.yml
 
+wait_for_resources:
+  extends: [.integration_tests_rules]
+  image: python:3.10
+  timeout: 7 days
+  tags:
+    - arch/amd64
+    - env/prod
+    - origin/jet-fleet
+    - owner/jet-core
+    - purpose/utility
+    - team/megatron
+  script:
+    - env
+    - pip install --no-cache-dir python-gitlab click
+    - export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}
+    - export GITLAB_ENDPOINT
+    - python tests/test_utils/python_scripts/wait_for_resources.py --pipeline-id $CI_PIPELINE_ID
+  rules:
+    - if: $INTEGRATION_TEST == "yes" && $CI_PIPELINE_SOURCE == "merge_request_event"
+      when: on_success
+    - when: never
+
 integration:configure:
   needs:
+    - job: wait_for_resources
+      optional: true
     - test:build_image
     - job: test:unit_tests_pyt(DEV)_mcore(latest)
       optional: true
@@ -78,6 +102,7 @@ integration:configure:
   needs:
     - integration:configure
     - test:build_image
+    - wait_for_resources
   extends: [.integration_tests_rules]
   trigger:
     include:
@@ -93,7 +118,6 @@ integration:configure:
     DASHBOARD_ENDPOINT: $DASHBOARD_ENDPOINT
     MCORE_MR_COMMIT: $MCORE_MR_COMMIT
     MCORE_BACKWARDS_COMMIT: $MCORE_BACKWARDS_COMMIT
-
   inherit:
     variables: true
 
diff --git a/.gitlab/stages/05.publish.yml b/.gitlab/stages/05.publish.yml
index f0d2a618..81125607 100644
--- a/.gitlab/stages/05.publish.yml
+++ b/.gitlab/stages/05.publish.yml
@@ -466,10 +466,17 @@ publish:docs:
     - owner/jet-core
     - purpose/utility
     - team/megatron
+  before_script:
+    - eval PUBLISH_COMMIT=$PUBLISH_COMMIT
+    - git fetch origin '+refs/merge-requests/*:refs/remotes/merge-requests/*'
+    - git fetch origin $PUBLISH_COMMIT
+    - git checkout $PUBLISH_COMMIT
   script:
     - cd ..
-    - rm -rf documentation && git clone --recursive https://gitlab-ci-token:${PROJECT_ACCESS_TOKEN_MCORE}@${GITLAB_ENDPOINT}/nemo-megatron-core-tme/documentation.git
+    - rm -rf documentation && git clone --recursive https://gitlab-ci-token:${PAT}@${GITLAB_ENDPOINT}/nemo-megatron-core-tme/documentation.git
     - cd documentation/megatron-lm
+    - git config --global user.email "mcore-bot@nvidia.com"
+    - git config --global user.name "Mcore Bot"
     - git fetch origin '+refs/merge-requests/*:refs/remotes/merge-requests/*'
     - git fetch origin $PUBLISH_COMMIT
     - git checkout $PUBLISH_COMMIT
@@ -479,6 +486,10 @@ publish:docs:
       git commit -m 'feat: Bump mcore'
 
     - git push
+  rules:
+    - if: '$CI_COMMIT_REF_PROTECTED == "true" && $CI_PIPELINE_SOURCE == "push"'
+      allow_failure: true
+    - when: never
 
 publish:upload_statistics:
   stage: publish
@@ -518,7 +529,7 @@ public:review_reminder:
   image: ${UTILITY_IMAGE}:${CI_PIPELINE_ID}
   script:
     - export GITLAB_ENDPOINT
-    - export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}
+    - export RO_API_TOKEN=${PAT}
     - export SLACK_WEBHOOK_URL=${SLACK_REMINDER_HOOK}
     - export SLACK_API_TOKEN=${SLACK_API_TOKEN}
     - python tests/test_utils/python_scripts/auto_reminder.py
@@ -530,6 +541,5 @@ public:review_reminder:
     - purpose/utility
     - team/megatron
   rules:
-    - if: $CI_COMMIT_BRANCH == "ko3n1g/ci/auto-reminder"
-    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH && $PUBLISH == "yes" && $PUBLISH_SCOPE == "review-reminder"
+    - if: $CI_COMMIT_BRANCH == "ci-review-reminder" && $PUBLISH == "yes" && $PUBLISH_SCOPE == "review-reminder"
     - when: never
diff --git a/examples/multimodal/mlp_converter.py b/examples/multimodal/mlp_converter.py
new file mode 100644
index 00000000..b105b93d
--- /dev/null
+++ b/examples/multimodal/mlp_converter.py
@@ -0,0 +1,75 @@
+# Copyright (c) 2024, FlagScale CORPORATION. All rights reserved.
+import argparse
+import os
+
+import torch
+
+
+def convert(input_path, output_path, tensor_parallel_size):
+    device = "cuda"
+
+    state_dict = torch.load(input_path, weights_only=False)
+
+    new_state_dicts = [{"model": dict()} for _ in range(tensor_parallel_size)]
+
+    for name, tensor in state_dict.items():
+
+        # Map parameter names to ones used in megatron.
+        new_name = ""
+        new_tensor = tensor
+        chunk_dim = None
+
+        # This is used for chunking some tensors to target tensor parallel size.
+        if name == "model.mm_projector.0.weight":
+            new_name = "encoder.linear_fc1.weight"
+            chunk_dim = 0
+        elif name == "model.mm_projector.0.bias":
+            new_name = "encoder.linear_fc1.bias"
+            chunk_dim = 0
+        elif name == "model.mm_projector.2.weight":
+            new_name = "encoder.linear_fc2.weight"
+            chunk_dim = 1
+        elif name == "model.mm_projector.2.bias":
+            new_name = "encoder.linear_fc2.bias"
+
+        assert new_name != "", f"unexpected name {name}"
+
+        if chunk_dim is None:
+            new_tensors = [new_tensor for _ in range(tensor_parallel_size)]
+        else:
+            new_tensors = torch.chunk(new_tensor, tensor_parallel_size, dim=chunk_dim)
+
+        for i in range(tensor_parallel_size):
+            # chunk() creates a view of a bigger tensor. clone() is used here to avoid excessive storage.
+            new_state_dicts[i]["model"][new_name] = new_tensors[i].clone()
+
+    for i in range(tensor_parallel_size):
+        output_path_tp = os.path.join(output_path, f"state_dict_tp_{i}.pt")
+        torch.save(new_state_dicts[i], output_path_tp)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="""
+Convert LLaVA MLP weights to megatron format.
+
+
+Example usage:
+python mlp_converter.py --input /some/input/folder/mm_projector.bin --output /some/output/folder --tensor-parallel-size 2
+""",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    parser.add_argument("--input", type=str, required=True, help="The mlp weights with hf format")
+    parser.add_argument(
+        "--output", type=str, required=True, help="output directory for megatron state dict file(s)"
+    )
+    parser.add_argument(
+        "--tensor-parallel-size", type=int, default=1, help="model tensor parallel size"
+    )
+
+    args = parser.parse_args()
+
+    convert(args.input, args.output, args.tensor_parallel_size)
+
+    print("done.")
diff --git a/megatron/__init__.py b/megatron/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/megatron/core/datasets/blended_dataset.py b/megatron/core/datasets/blended_dataset.py
index fccd0068..a1be03f4 100644
--- a/megatron/core/datasets/blended_dataset.py
+++ b/megatron/core/datasets/blended_dataset.py
@@ -13,7 +13,7 @@ import torch
 
 from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
 from megatron.core.datasets.megatron_dataset import MegatronDataset
-from megatron.core.datasets.utils import normalize
+from megatron.core.datasets.utils import normalize, is_built_on_zero_rank
 from megatron.core.utils import log_single_rank
 
 logger = logging.getLogger(__name__)
@@ -122,7 +122,7 @@ class BlendedDataset(torch.utils.data.Dataset):
         else:
             cache_hit = False
 
-        if not path_to_cache or (not cache_hit and torch.distributed.get_rank() == 0):
+        if not path_to_cache or (not cache_hit and is_built_on_zero_rank()):
             log_single_rank(
                 logger, logging.INFO, f"Build and save the {type(self).__name__} indices"
             )
diff --git a/megatron/core/datasets/blended_megatron_dataset_builder.py b/megatron/core/datasets/blended_megatron_dataset_builder.py
index af9375e0..cc9f13cf 100644
--- a/megatron/core/datasets/blended_megatron_dataset_builder.py
+++ b/megatron/core/datasets/blended_megatron_dataset_builder.py
@@ -11,7 +11,7 @@ import torch
 from megatron.core.datasets.blended_dataset import BlendedDataset
 from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
 from megatron.core.datasets.megatron_dataset import LowLevelDataset, MegatronDataset
-from megatron.core.datasets.utils import Split, normalize
+from megatron.core.datasets.utils import Split, normalize, is_built_on_zero_rank
 from megatron.core.parallel_state import get_virtual_pipeline_model_parallel_rank
 from megatron.core.utils import log_single_rank
 
@@ -359,7 +359,7 @@ class BlendedMegatronDatasetBuilder(object):
         if torch.distributed.is_initialized():
             rank = torch.distributed.get_rank()
             # First, build on rank 0
-            if rank == 0:
+            if is_built_on_zero_rank():
                 num_workers = num_dataset_builder_threads
                 if num_workers > 1:
                     # since only rank 0 is running, scale up the thread count
@@ -374,7 +374,7 @@ class BlendedMegatronDatasetBuilder(object):
             torch.distributed.barrier()
 
             # Then, build on other ranks; guaranteed to be data_cache hit
-            if rank != 0:
+            if not is_built_on_zero_rank():
                 _threading_helper(
                     megatron_datasets,
                     num_dataset_builder_threads,
@@ -491,7 +491,7 @@ class BlendedMegatronDatasetBuilder(object):
             dataset = None
 
             # First, build on rank 0
-            if rank == 0 and is_built_on_rank():
+            if is_built_on_zero_rank() and is_built_on_rank():
                 try:
                     dataset = cls(*args)
                 except OSError as err:
@@ -507,7 +507,7 @@ class BlendedMegatronDatasetBuilder(object):
                 torch.distributed.barrier()
 
             # After, build on other ranks
-            if rank != 0 and is_built_on_rank():
+            if not is_built_on_zero_rank() and is_built_on_rank():
                 dataset = cls(*args)
 
             return dataset
diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py
index cf805491..5582eea1 100644
--- a/megatron/core/datasets/gpt_dataset.py
+++ b/megatron/core/datasets/gpt_dataset.py
@@ -14,7 +14,7 @@ from megatron.core.datasets.indexed_dataset import IndexedDataset
 from megatron.core.datasets.megatron_dataset import MegatronDataset
 from megatron.core.datasets.megatron_tokenizer import MegatronTokenizer
 from megatron.core.datasets.object_storage_utils import ObjectStorageConfig, is_object_storage_path
-from megatron.core.datasets.utils import Split
+from megatron.core.datasets.utils import Split, is_built_on_zero_rank
 from megatron.core.utils import log_single_rank
 
 logger = logging.getLogger(__name__)
@@ -354,7 +354,7 @@ class GPTDataset(MegatronDataset):
 
         if not path_to_cache or (
             not cache_hit
-            and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0)
+            and (not torch.distributed.is_initialized() or is_built_on_zero_rank())
         ):
 
             log_single_rank(
diff --git a/megatron/core/datasets/utils.py b/megatron/core/datasets/utils.py
index 8d887d4a..058f035c 100644
--- a/megatron/core/datasets/utils.py
+++ b/megatron/core/datasets/utils.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 
+import os
 import logging
 from enum import Enum
 from typing import List, Optional, Tuple
@@ -85,3 +86,26 @@ def get_blend_from_list(
     prefix_per_dataset = [rppd.strip() for rppd in raw_prefix_per_dataset]
 
     return prefix_per_dataset, weight_per_dataset
+
+
+def is_built_on_zero_rank():
+    """
+    Determines if the current distributed rank is the one responsible for building datasets.
+
+    Returns:
+        bool: True if the current rank is responsible for building resources, False otherwise.
+    """
+    from megatron.training import get_args
+    args = get_args()
+
+    is_built = False
+    if not args.no_shared_fs \
+        and torch.distributed.get_rank() == 0:
+        is_built = True
+    elif args.no_shared_fs \
+        and int(os.environ["LOCAL_RANK"]) == 0:
+        is_built = True
+    else:
+        is_built = False
+
+    return is_built
\ No newline at end of file
diff --git a/megatron/core/dist_checkpointing/exchange_utils.py b/megatron/core/dist_checkpointing/exchange_utils.py
index 8486c7ef..9fbc0158 100644
--- a/megatron/core/dist_checkpointing/exchange_utils.py
+++ b/megatron/core/dist_checkpointing/exchange_utils.py
@@ -62,7 +62,7 @@ class ShardDistribution(NamedTuple):
 def _shard_size(sh_ten: ShardedTensor):
     """Returns size in bytes of a given sharded tensor."""
     if sh_ten.flattened_range is None:
-        numel = np.product(sh_ten.local_shape)
+        numel = np.prod(sh_ten.local_shape)
     else:
         numel = sh_ten.flattened_range.stop - sh_ten.flattened_range.start
     return numel * torch._utils._element_size(sh_ten.dtype)
diff --git a/megatron/core/dist_checkpointing/mapping.py b/megatron/core/dist_checkpointing/mapping.py
index 156702b2..6028b6b4 100644
--- a/megatron/core/dist_checkpointing/mapping.py
+++ b/megatron/core/dist_checkpointing/mapping.py
@@ -204,7 +204,7 @@ class ShardedTensor(ShardedBase):
             )
 
         # TODO: np.unravel_index?
-        mask = np.zeros(np.product(self.local_shape), dtype=bool)
+        mask = np.zeros(np.prod(self.local_shape), dtype=bool)
         mask[self.flattened_range] = True
         return np.nonzero(mask.reshape(self.local_shape))
 
diff --git a/megatron/core/dist_checkpointing/serialization.py b/megatron/core/dist_checkpointing/serialization.py
index 8cc5597e..e534590f 100644
--- a/megatron/core/dist_checkpointing/serialization.py
+++ b/megatron/core/dist_checkpointing/serialization.py
@@ -8,6 +8,7 @@ Additionally, `load` expects the sharded state dict argument as a guidance for
 loading the sharded tensors.
 """
 
+import os
 import logging
 from pathlib import Path
 from typing import Callable, Dict, Optional, Set, Tuple, Union
diff --git a/megatron/core/dist_checkpointing/strategies/filesystem_async.py b/megatron/core/dist_checkpointing/strategies/filesystem_async.py
index 2af90a65..30a52212 100644
--- a/megatron/core/dist_checkpointing/strategies/filesystem_async.py
+++ b/megatron/core/dist_checkpointing/strategies/filesystem_async.py
@@ -5,21 +5,23 @@ import dataclasses
 import logging
 import os
 import queue
+import pickle
 from functools import partial
 from heapq import heappop, heappush
 from itertools import chain
 from operator import itemgetter
 from pathlib import Path
 from time import time
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union, cast
 
 import psutil
 import torch
 from torch import multiprocessing as mp
 from torch.distributed.checkpoint import FileSystemWriter
-from torch.distributed.checkpoint.filesystem import DEFAULT_SUFFIX, _StoragePrefix, _write_item
+from torch.distributed.checkpoint.filesystem import DEFAULT_SUFFIX, _StoragePrefix, _write_item, _metadata_fn
 from torch.distributed.checkpoint.planner import SavePlan, SavePlanner, WriteItem, WriteItemType
 from torch.distributed.checkpoint.storage import WriteResult
+from torch.distributed.checkpoint.metadata import Metadata
 from torch.futures import Future
 
 from .async_utils import _disable_gc
@@ -30,6 +32,40 @@ WriteBucket = Tuple[Path, str, Tuple[list, list]]  # represents writes to a sing
 
 _results_queue = None
 
+_GLOBAL_PREVIOUS_METADATA = None
+
+_GLOBAL_PREVIOUS_COUNT = 0
+
+
+def get_previous_metadata():
+    """
+    Get the metadata from the previous save.
+    """
+    return _GLOBAL_PREVIOUS_METADATA
+
+
+def set_previous_metadata(metadata):
+    """
+    Set the metadata from the previous save.
+    """
+    global _GLOBAL_PREVIOUS_METADATA
+    _GLOBAL_PREVIOUS_METADATA = metadata
+
+
+def get_previous_count():
+    """
+    Get the count from the previous save.
+    """
+    return _GLOBAL_PREVIOUS_COUNT
+
+
+def set_previous_count(count):
+    """
+    Set the count from the previous save.
+    """
+    global _GLOBAL_PREVIOUS_COUNT
+    _GLOBAL_PREVIOUS_COUNT = count
+
 
 def _get_write_results_queue():
     global _results_queue
@@ -75,6 +111,13 @@ class FileSystemWriterAsync(FileSystemWriter):
         self.results_queue: Optional[mp.Queue] = None
         self.separation_hint = separation_hint
 
+        # Get the value from the environment variable if it exists, otherwise default to False
+        self.single_file_per_tensor_ckpt = os.getenv('FS_SFPT_CKPT_SAVE', 'False').lower() in (
+            'true',
+            '1',
+            't',
+        )
+
     def prepare_write_data(self, plan: SavePlan, planner: SavePlanner) -> None:
         """
         First stage of async saving. Copy data to CPU and plan the local saving.
@@ -99,12 +142,17 @@ class FileSystemWriterAsync(FileSystemWriter):
         start = time()
         # move tensors from GPU to CPU before starting async writing
         # We do D2H synchronously for now
-        file_count = 0
+        if not self.single_file_per_tensor_ckpt:
+            file_count = 0
+        else:
+            file_count = get_previous_count()
 
         def gen_file(prefix=""):
             nonlocal file_count
             file_name = f"{prefix}{storage_plan.prefix}{file_count}{DEFAULT_SUFFIX}"
             file_count += 1
+            if self.single_file_per_tensor_ckpt:
+                set_previous_count(file_count)
             return file_name
 
         def _clone_if_needed(ten: torch.Tensor):
@@ -392,6 +440,48 @@ class FileSystemWriterAsync(FileSystemWriter):
             local_plan, storage_data=_StoragePrefix(f"__{torch.distributed.get_rank()}_")
         )
 
+    def finish(self, metadata: Metadata, results: List[List[WriteResult]]) -> None:
+        # Modify based on the original implementation from torch.distributed.checkpoint.filesystem.FileSystemWriter
+        # https://github.com/pytorch/pytorch/blob/625c24a7f98a645b6f8758a01d7163a842582ce0/torch/distributed/checkpoint/filesystem.py#L574
+
+        if not self.single_file_per_tensor_ckpt:
+            storage_md = {}
+        else:
+            if get_previous_count() == 1:
+                storage_md = {}
+            else:
+                # Get the metadata from the previous save
+                prev_metadata = get_previous_metadata()
+                prev_metadata.state_dict_metadata.update(metadata.state_dict_metadata)
+                metadata = prev_metadata
+                storage_md = metadata.storage_data
+
+        for wr_list in results:
+            storage_md.update({wr.index: wr.storage_data for wr in wr_list})
+        metadata.storage_data = storage_md
+
+        if not self.single_file_per_tensor_ckpt or get_previous_count() == 1:
+            metadata.storage_meta = self.storage_meta()
+
+        tmp_path = cast(Path, self.fs.concat_path(self.path, f"{_metadata_fn}.tmp"))
+        with self.fs.create_stream(tmp_path, "wb") as metadata_file:
+            pickle.dump(metadata, metadata_file)
+            if self.sync_files:
+                try:
+                    os.fsync(metadata_file.fileno())
+                except AttributeError:
+                    os.sync()
+
+        # delete in-case other checkpoints were present.
+        if self.fs.exists(self.metadata_path):
+            self.fs.rm_file(self.metadata_path)
+
+        self.fs.rename(tmp_path, self.metadata_path)
+
+        # Store the metadata for the next save
+        if self.single_file_per_tensor_ckpt:
+            set_previous_metadata(metadata)
+
 
 def _split_by_size_and_type(bins: int, items: List[WriteItem]) -> List[List[WriteItem]]:
     """
diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py
index 06c262a8..12941255 100644
--- a/megatron/core/dist_checkpointing/strategies/torch.py
+++ b/megatron/core/dist_checkpointing/strategies/torch.py
@@ -855,6 +855,13 @@ class TorchDistLoadShardedStrategy(LoadShardedStrategy):
 
         Returns: loaded state dict
         """
+        # Get the value from the environment variable if it exists, otherwise default to True
+        single_file_per_tensor_ckpt = os.getenv('FS_SFPT_CKPT_LOAD', 'False').lower() in (
+            'true',
+            '1',
+            't',
+        )
+
         # Apply N-D tensors resharding
         reformulation_metadata = get_reformulation_metadata(sharded_state_dict, checkpoint_dir)
         sharded_state_dict, formulation_restore_data = apply_nd_flattened_tensors_reformulation(
@@ -889,14 +896,24 @@ class TorchDistLoadShardedStrategy(LoadShardedStrategy):
         )
         # Load PyT Distributed format
         fsr = CachedMetadataFileSystemReader(checkpoint_dir)
-        checkpoint.load_state_dict(
-            pyt_state_dict,
-            fsr,
-            planner=MCoreLoadPlanner(
-                shapes_validation_sharded_tensors=flexible_shape_sharded_tensors,
+        if not single_file_per_tensor_ckpt:
+            checkpoint.load_state_dict(
+                pyt_state_dict,
+                fsr,
+                planner=MCoreLoadPlanner(
+                    shapes_validation_sharded_tensors=flexible_shape_sharded_tensors,
                 allow_shape_mismatch_sharded_tensors=allow_shape_mismatch_sharded_tensors,
-            ),
-        )
+                ),
+            )
+        else:
+            checkpoint.load_state_dict(
+                pyt_state_dict,
+                fsr,
+                planner=MCoreLoadPlanner(
+                    shapes_validation_sharded_tensors=flexible_shape_sharded_tensors,
+                    allow_partial_load=True,
+                ),
+            )
 
         self.cached_global_metadata = (
             fsr.read_metadata()
@@ -910,6 +927,13 @@ class TorchDistLoadShardedStrategy(LoadShardedStrategy):
             k: v if not isinstance(v, TorchShardedTensor) else _unwrap_pyt_sharded_tensor(v)
             for k, v in pyt_state_dict.items()
         }
+
+        if single_file_per_tensor_ckpt:
+            mcore_state_dict = {
+                k: [None] if (not isinstance(v, list) and "_extra_state" in k) else v
+                for k, v in mcore_state_dict.items()
+            }
+
         mcore_state_dict = _replace_sharded_keys_with_state_dict_keys(
             mcore_state_dict, flat_mapping, rename_mapping
         )
diff --git a/megatron/core/dist_checkpointing/validation.py b/megatron/core/dist_checkpointing/validation.py
index 546ec354..f4572f11 100644
--- a/megatron/core/dist_checkpointing/validation.py
+++ b/megatron/core/dist_checkpointing/validation.py
@@ -494,7 +494,7 @@ def _validate_sharding_for_key_flattened(tensors_by_shard):
         all_slices.append((sharding.flattened_range.start, sharding.flattened_range.stop))
 
     starts, stops = map(np.asarray, zip(*sorted(all_slices)))
-    expected_size = np.product(local_shape)
+    expected_size = np.prod(local_shape)
     if starts[0] != 0 or stops[-1] != expected_size or not np.all(starts[1:] == stops[:-1]):
         raise CheckpointingException(
             f'Flattened ranges dont cover the whole shard {tensors_by_shard[0]} of size {expected_size}. Ranges: {(starts, stops)}'
diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py
index 9a28b078..58f752ee 100644
--- a/megatron/core/distributed/finalize_model_grads.py
+++ b/megatron/core/distributed/finalize_model_grads.py
@@ -25,6 +25,18 @@ def _get_main_grad_attr(param: torch.nn.Parameter, use_custom_fsdp: bool = False
         return "main_grad"
     return "grad"
 
+def get_device_type_for_comm(model_parallel_group=None):
+    ''''Copy from flagscale/train/hetero/p2p_communication.py'''
+    device = 'cuda'
+    # "cpu:gloo": gloo only supports cpu tensor.
+    # "gloo" & "cpu:gloo,cuda:gloo": gloo supports both cpu and cuda tensor.
+    if isinstance(model_parallel_group, list):
+        if 'cpu:gloo' == torch.distributed.get_backend(model_parallel_group[0]):
+            device = 'cpu'
+    else:
+        if 'cpu:gloo' == torch.distributed.get_backend(model_parallel_group):
+            device = 'cpu'
+    return device
 
 def _unshard_if_dtensor(tensor: Union[torch.Tensor, "DTensor"]) -> torch.Tensor:
     """
@@ -126,8 +138,14 @@ def _allreduce_word_embedding_grads(model: List[torch.nn.Module], config: Transf
 
     if (
         parallel_state.is_rank_in_embedding_group(ignore_virtual=True)
-        and torch.distributed.get_world_size(parallel_state.get_embedding_group()) > 1
     ):
+        embed_group = parallel_state.get_embedding_group()
+        if not isinstance(embed_group, list):
+            embed_group = [embed_group]
+    else:
+        return
+
+    if (torch.distributed.get_world_size(embed_group[0]) > 1):
         if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
             model_module = model[0]
         elif parallel_state.is_pipeline_last_stage(ignore_virtual=True):
@@ -136,6 +154,7 @@ def _allreduce_word_embedding_grads(model: List[torch.nn.Module], config: Transf
             model_module = model[0]
 
         ddp_config = model_module.ddp_config
+        use_dist_opt = ddp_config.use_distributed_optimizer
         model_module = get_attr_wrapped_model(model_module, 'pre_process', return_model_obj=True)
 
         # If share_embeddings_and_output_weights is True, we need to maintain duplicated
@@ -147,7 +166,37 @@ def _allreduce_word_embedding_grads(model: List[torch.nn.Module], config: Transf
             grad_attr = _get_main_grad_attr(weight, ddp_config.use_custom_fsdp)
             orig_grad = getattr(weight, grad_attr)
             grad = _unshard_if_dtensor(orig_grad)
-            torch.distributed.all_reduce(grad, group=parallel_state.get_embedding_group())
+            com_device = get_device_type_for_comm(embed_group)
+            if com_device == "cpu":
+                grad = grad.cpu()
+            if use_dist_opt:
+                if config.use_partial_reduce_for_shared_embedding:
+                    dp_world_size = parallel_state.get_data_parallel_world_size()
+                    dp_rank = parallel_state.get_data_parallel_rank()
+                    assert grad.shape[0] % dp_world_size == 0, f"grad shape: {grad.shape[0]}, dp_world_size: {dp_world_size}"
+                    per_partion_size = grad.shape[0] // dp_world_size
+                    if len(embed_group) == 1:
+                        offset = per_partion_size * dp_rank
+                        torch.distributed.all_reduce(grad[offset:offset+per_partion_size, :], group=embed_group[0])
+                    else:
+                        group_idx = 0
+                        per_partion_size = per_partion_size // len(embed_group)
+                        for group in embed_group:
+                            offset = per_partion_size * (dp_rank * len(embed_group) + group_idx)
+                            torch.distributed.all_reduce(grad[offset : offset + per_partion_size, :], group=group)
+                            group_idx += 1
+                else: # megartron default method
+                    torch.distributed.all_reduce(grad, group=embed_group[0])
+            else:
+                if len(embed_group) == 1: # megartron default method
+                    torch.distributed.all_reduce(grad, group=embed_group[0])
+                else:
+                    original_grad_data = grad.clone().detach().data
+                    for group in embed_group:
+                        grad.data.copy_(original_grad_data)
+                        torch.distributed.all_reduce(grad, group=group)
+            if grad.device == torch.device('cpu'):
+                grad.to(torch.cuda.current_device())
             setattr(weight, grad_attr, _reshard_if_dtensor(grad, orig_grad))
 
 
@@ -310,6 +359,11 @@ def finalize_model_grads(model: List[torch.nn.Module], num_tokens: Optional[torc
         last_rank = parallel_state.get_pipeline_model_parallel_last_rank()
         pp_group = parallel_state.get_pipeline_model_parallel_group()
 
+        # NOTE: This is a hack to support multiple pipeline parallel groups. The origin
+        #       parallel_state.get_pipeline_model_parallel_last_rank() only supports a single
+        if isinstance(pp_group, list):
+            last_rank = [parallel_state.get_pipeline_model_parallel_last_rank(g) for g in pp_group]
+
         if not isinstance(last_rank, list):
             assert not isinstance(last_rank, list)
             last_rank = [last_rank]
@@ -317,12 +371,18 @@ def finalize_model_grads(model: List[torch.nn.Module], num_tokens: Optional[torc
             pp_group = [pp_group]
 
         # need to do a broadcast for every pp group, even though num_tokens should be the same.
+        if "cpu:gloo" == pp_group[0].name():
+            num_tokens = num_tokens.cpu()
+
         num_tokens_list = []
         for lr, group in zip(last_rank, pp_group):
             torch.distributed.broadcast(num_tokens, src=lr, group=group)
             num_tokens_list.append(torch.clone(num_tokens))
         assert all(x.item() == num_tokens_list[0] for x in num_tokens_list)
 
+        if num_tokens.device == torch.device('cpu'):
+            num_tokens = num_tokens.cuda()
+
         # all-reduce across DP ranks.
         torch.distributed.all_reduce(num_tokens, group=parallel_state.get_data_parallel_group())
         for model_chunk in model:
diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index b60e5056..1404e7dd 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -23,6 +23,7 @@ from megatron.core.parallel_state import (
     get_expert_model_parallel_world_size,
     get_hierarchical_context_parallel_groups,
     get_tensor_model_parallel_group,
+    get_tensor_model_parallel_world_size,
 )
 from megatron.core.process_groups_config import ModelCommProcessGroups
 from megatron.core.tensor_parallel.layers import (
@@ -372,7 +373,7 @@ class TELayerNormColumnParallelLinear(te.pytorch.LayerNormLinear):
             sequence_parallel=self.config.sequence_parallel,
             fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
             tp_group=tp_group if torch.distributed.is_initialized() else None,
-            tp_size=self.config.tensor_model_parallel_size,
+            tp_size=get_tensor_model_parallel_world_size(),
             get_rng_state_tracker=(
                 get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
             ),
@@ -765,7 +766,7 @@ class TEDotProductAttention(te.pytorch.DotProductAttention):
             ),
             attn_mask_type=attn_mask_type.name,
             sequence_parallel=self.config.sequence_parallel,
-            tp_size=self.config.tensor_model_parallel_size,
+            tp_size=get_tensor_model_parallel_world_size(),
             get_rng_state_tracker=(
                 get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
             ),
@@ -1034,7 +1035,7 @@ if is_te_min_version("1.9.0.dev0"):
                 return pickle.loads(state.detach().cpu().numpy().tobytes())
             elif isinstance(state, io.BytesIO):
                 state.seek(0)
-                return torch.load(state, map_location="cuda")
+                return torch.load(state, map_location="cuda", weights_only=False)
             else:
                 raise RuntimeError("Unsupported checkpoint format.")
 
diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py
index e64ad37a..24e2d678 100644
--- a/megatron/core/model_parallel_config.py
+++ b/megatron/core/model_parallel_config.py
@@ -343,6 +343,16 @@ class ModelParallelConfig:
        the user adds a level 1 timer that is not called by all ranks.
     """
 
+    ###################
+    # Heterogeneous Training
+    ###################
+    enable_hetero: str = None
+    """Enable the heterogeneous training."""
+
+    hetero_pipeline_layer_split: list = None
+    """A list of lists, each sublist contains numbers of layers to be processed in the corresponding pipeline stages for one device type."""
+
+
     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
diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py
index b777de3a..9412be03 100644
--- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py
+++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py
@@ -223,7 +223,7 @@ class RotaryEmbedding(nn.Module):
                 rotary_seq_len = transformer_input.size(0)
 
             if transformer_config.sequence_parallel:
-                rotary_seq_len *= transformer_config.tensor_model_parallel_size
+                rotary_seq_len *= parallel_state.get_tensor_model_parallel_world_size()
 
         rotary_seq_len *= transformer_config.context_parallel_size
 
diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index 4022ac57..5e26bc99 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -162,9 +162,26 @@ class LanguageModule(MegatronModule):
             if parallel_state.is_rank_in_embedding_group():
                 weight = self.shared_embedding_or_output_weight()
                 weight.data = weight.data.cuda()
-                torch.distributed.all_reduce(
-                    weight.data, group=parallel_state.get_embedding_group()
-                )
+                embedding_group = parallel_state.get_embedding_group()
+                if not isinstance(embedding_group, list):
+                    torch.distributed.all_reduce(
+                        weight.data, group=parallel_state.get_embedding_group()
+                    )
+                else: # for multiple embedding groups in heterogeneous mode
+                    with torch.no_grad():
+                        original_dtype = weight.dtype
+                        if original_dtype == torch.bfloat16: # gloo backend doesn't support bfloat16
+                            weight = weight.to(torch.float32)
+                        if torch.distributed.get_backend(group=embedding_group[0]) == 'cpu:gloo':
+                            weight.data = weight.data.cpu()
+                        original_weight = weight.clone().detach().data
+                        for group in embedding_group:
+                            weight.data.copy_(original_weight)
+                            torch.distributed.all_reduce(weight.data, group=group)
+                        if original_dtype == torch.bfloat16:
+                            weight = weight.to(original_dtype)
+                        if weight.device == torch.device('cpu'):
+                            weight.data = weight.data.cuda()
 
         elif not getattr(LanguageModule, "embedding_warning_printed", False):
             logging.getLogger(__name__).warning(
diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
old mode 100755
new mode 100644
diff --git a/megatron/core/models/multimodal/llava_model.py b/megatron/core/models/multimodal/llava_model.py
index 44e48238..abc38c77 100644
--- a/megatron/core/models/multimodal/llava_model.py
+++ b/megatron/core/models/multimodal/llava_model.py
@@ -328,9 +328,16 @@ class LLaVAModel(MegatronModule):
                     f"vision_projection.{name}"
                     for name in self.vision_projection.state_dict().keys()
                 ]
+                vision_extra_state_param_names = []
+                for name in self.vision_model.state_dict().keys():
+                    if "_extra_state" in name:
+                        vision_extra_state_param_names.append(f"vision_model.{name}")
                 self.vision_projection.register_load_state_dict_post_hook(
                     partial(_load_state_dict_hook_ignore_param_names, vision_projection_param_names)
                 )
+                self.vision_model.register_load_state_dict_post_hook(
+                    partial(_load_state_dict_hook_ignore_param_names, vision_extra_state_param_names)
+                )
 
             self.vision_projection.register_load_state_dict_post_hook(
                 _load_state_dict_hook_ignore_extra_state
diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index 2608fde8..a192e681 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -24,6 +24,13 @@ except ImportError:
         # See https://github.com/NVIDIA/apex/blob/7b73b12361068a10b0f44844534613f252a5ea75/apex/optimizers/fused_adam.py#L16.
         from torch.optim import AdamW as Adam, SGD
 
+import os
+if not os.environ.get('DISABLE_QADAM'):
+    import sys
+    so_dir_path = "/workspace/"
+    sys.path.append(so_dir_path)
+    from qadam import FusedAdam as Adam
+
 from megatron.core import mpu
 from megatron.core.optimizer.cpu_offloading.hybrid_optimizer import HybridDeviceOptimizer
 
@@ -464,7 +471,9 @@ def get_megatron_optimizer(
     else:
         all_dense_model_chunks = [model_chunks]
         overlap_param_gather_with_optimizer_step_flags = [False]
-    model_parallel_rank = torch.distributed.get_rank(mpu.get_model_parallel_group())
+    mp_group = mpu.get_model_parallel_group()
+    mp_group = [mp_group] if not isinstance(mp_group, list) else mp_group
+    model_parallel_rank = torch.distributed.get_rank(mp_group[0])
 
     if torch.distributed.get_world_size(
         mpu.get_data_parallel_group(with_context_parallel=True, partial_data_parallel=False)
@@ -568,9 +577,14 @@ def get_megatron_optimizer(
         buffer_name='expert_parallel_buffers',
     )
     if len(moe_param_groups) > 0:
-        model_parallel_rank = torch.distributed.get_rank(
-            mpu.get_expert_tensor_model_pipeline_parallel_group()
-        )
+        expert_mp_group = mpu.get_expert_tensor_model_pipeline_parallel_group()
+        if not isinstance(expert_mp_group, list):
+            model_parallel_rank = torch.distributed.get_rank(
+                mpu.get_expert_tensor_model_pipeline_parallel_group()
+            )
+        else:
+            model_parallel_rank = torch.distributed.get_rank(expert_mp_group[0])
+
         # Pass Gloo process groups into optimizer only if needed.
         if use_gloo_process_groups:
             data_parallel_group_gloo = mpu.get_expert_data_parallel_group_gloo()
diff --git a/megatron/core/optimizer/clip_grads.py b/megatron/core/optimizer/clip_grads.py
index 0f33f919..e2002c49 100644
--- a/megatron/core/optimizer/clip_grads.py
+++ b/megatron/core/optimizer/clip_grads.py
@@ -47,6 +47,7 @@ from ..tensor_parallel import param_is_not_tensor_parallel_duplicate
 from ..transformer.module import param_is_not_shared
 from ..utils import get_data_parallel_group_if_dtensor, to_local_if_dtensor
 
+from flagscale.train.hetero.p2p_communication import get_device_type_for_comm
 
 def get_grad_norm_fp32(
     grads_for_norm: Union[List[torch.Tensor], torch.Tensor],
@@ -93,9 +94,20 @@ def get_grad_norm_fp32(
             torch.distributed.all_reduce(
                 total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=data_parallel_group
             )
-        torch.distributed.all_reduce(
-            total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=grad_stats_parallel_group
-        )
+
+        # Take max across all model-parallel GPUs.
+        # For cpu comminication
+        tensor_device = get_device_type_for_comm(grad_stats_parallel_group)
+        total_norm_cuda.to(tensor_device)
+        if isinstance(grad_stats_parallel_group, list):
+            for group in grad_stats_parallel_group:
+                torch.distributed.all_reduce(
+                    total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=group
+                )
+        else:
+            torch.distributed.all_reduce(
+                total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=grad_stats_parallel_group
+            )
         total_norm = total_norm_cuda[0].item()
 
     else:
@@ -127,9 +139,22 @@ def get_grad_norm_fp32(
             torch.distributed.all_reduce(
                 total_norm, op=torch.distributed.ReduceOp.SUM, group=data_parallel_group
             )
-        torch.distributed.all_reduce(
-            total_norm, op=torch.distributed.ReduceOp.SUM, group=grad_stats_parallel_group
-        )
+        # Sum across all model-parallel GPUs.
+        # For cpu comminication
+        tensor_device = get_device_type_for_comm(grad_stats_parallel_group)
+        total_norm = total_norm.to(tensor_device)
+        if isinstance(grad_stats_parallel_group, list):
+            original_total_norm = total_norm.clone().detach()
+            for mp_group in grad_stats_parallel_group:
+                total_norm.data = original_total_norm.data.clone()
+                total_norm = total_norm.to(tensor_device)
+                torch.distributed.all_reduce(
+                    total_norm, op=torch.distributed.ReduceOp.SUM, group=mp_group
+                )
+        else:
+            torch.distributed.all_reduce(
+                total_norm, op=torch.distributed.ReduceOp.SUM, group=grad_stats_parallel_group
+            )
         total_norm = total_norm.item() ** (1.0 / norm_type)
 
     return total_norm
@@ -223,9 +248,21 @@ def count_zeros_fp32(
             total_num_zeros, op=torch.distributed.ReduceOp.SUM, group=data_parallel_group
         )
     # Sum across all model-parallel GPUs.
-    torch.distributed.all_reduce(
-        total_num_zeros, op=torch.distributed.ReduceOp.SUM, group=grad_stats_parallel_group
-    )
+    comm_device = get_device_type_for_comm(grad_stats_parallel_group)
+    if comm_device == "cpu":
+        total_num_zeros = total_num_zeros.cpu()
+
+    if isinstance(grad_stats_parallel_group, list):
+        original_total_num_zeros = total_num_zeros.clone().detach()
+        for group in grad_stats_parallel_group:
+            total_num_zeros.data = original_total_num_zeros.data.clone()
+            torch.distributed.all_reduce(
+                total_num_zeros, op=torch.distributed.ReduceOp.SUM, group=group
+            )
+    else:
+        torch.distributed.all_reduce(
+            total_num_zeros, op=torch.distributed.ReduceOp.SUM, group=grad_stats_parallel_group
+        )
 
     total_num_zeros = total_num_zeros.item()
 
diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py
index 8bb3bc27..4715958c 100644
--- a/megatron/core/optimizer/distrib_optimizer.py
+++ b/megatron/core/optimizer/distrib_optimizer.py
@@ -7,6 +7,7 @@ import itertools
 from dataclasses import replace
 from logging import getLogger
 from typing import Callable, Dict, List, Optional, Tuple
+import os
 
 import torch
 
@@ -21,6 +22,16 @@ except ImportError:
 
         HAVE_APEX_OR_TE = False
 
+try:
+    import sys
+    so_dir_path = os.getcwd()
+    sys.path.append(so_dir_path)
+    from qadam import FusedAdam as Adam
+except ImportError:
+    warnings.warn(
+        f'qadam are not installed. Falling back.'
+    )
+
 from megatron.core.optimizer.cpu_offloading import HybridDeviceOptimizer
 
 from .. import tensor_parallel
diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py
index 5d2ec82f..b88f520b 100644
--- a/megatron/core/optimizer/optimizer.py
+++ b/megatron/core/optimizer/optimizer.py
@@ -393,12 +393,23 @@ class MixedPrecisionOptimizer(MegatronOptimizer):
             )
 
         # Update across all model parallel instances.
-        torch.distributed.all_reduce(
-            self.found_inf,
-            op=torch.distributed.ReduceOp.MAX,
-            group=self.get_grad_stats_parallel_group(),
-        )
-
+        groups = self.get_grad_stats_parallel_group()
+        if isinstance(groups, list):
+            if "cpu:gloo" == torch.distributed.get_backend(groups[0]):
+                self.found_inf = self.found_inf.cpu()
+        else:
+            if "cpu:gloo" == torch.distributed.get_backend(groups):
+                self.found_inf = self.found_inf.cpu()
+        if not isinstance(groups, list):
+            groups = [groups]
+        for group in groups:
+            torch.distributed.all_reduce(
+                self.found_inf,
+                op=torch.distributed.ReduceOp.MAX,
+                group=group
+            )
+        if self.found_inf.device != torch.device('cuda'):
+            self.found_inf = self.found_inf.cuda()
         # Check for nan.
         found_inf_flag = self.found_inf.item() > 0
 
@@ -1193,7 +1204,7 @@ class ChainedOptimizer(MegatronOptimizer):
 
             # Lazy loading checkpoint, state dict is needed only when DP rank = 0.
             if torch.distributed.get_rank(optimizer.data_parallel_group) == 0 and states is None:
-                states = torch.load(filename)
+                states = torch.load(filename, weights_only=False)
 
             state_dict = states[idx] if states else None
             optimizer.load_parameter_state_from_dp_zero(
diff --git a/megatron/core/optimizer_param_scheduler.py b/megatron/core/optimizer_param_scheduler.py
index 43c106f4..e14cc322 100644
--- a/megatron/core/optimizer_param_scheduler.py
+++ b/megatron/core/optimizer_param_scheduler.py
@@ -53,6 +53,7 @@ class OptimizerParamScheduler:
         override_opt_param_scheduler: Optional[bool] = False,
         wsd_decay_steps: Optional[int] = None,
         lr_wsd_decay_style: Optional[str] = None,
+        stablelm2_scheduler_config=None,
     ) -> None:
 
         # Class values.
@@ -91,6 +92,15 @@ class OptimizerParamScheduler:
                 'both override and ' 'use-checkpoint are set.'
             )
 
+        self.stablelm2_scheduler_config = stablelm2_scheduler_config
+        if self.stablelm2_scheduler_config is not None:
+          ## absolute samples
+          self.stablelm2_scheduler_config.rsqrt_samples += \
+              self.stablelm2_scheduler_config.cosine_samples
+          ## N of consine
+          if self.stablelm2_scheduler_config.cosine_period_samples == 0:
+            self.stablelm2_scheduler_config.cosine_period_samples = self.lr_decay_steps
+
         # Set the learning rate
         self.step(0)
         log_single_rank(logger, logging.INFO, f"> learning rate decay style: {self.lr_decay_style}")
@@ -150,6 +160,62 @@ class OptimizerParamScheduler:
             lr = max_lr * warmup_steps**0.5 / (num_steps**0.5)
             return max(min_lr, lr)
 
+        # stablelm2 scheduler of multiple stages
+        if self.stablelm2_scheduler_config is not None:
+            log_single_rank(logger, logging.INFO, f"> stablelm2_scheduler_config: {self.stablelm2_scheduler_config}")
+            if self.num_steps <= self.stablelm2_scheduler_config.cosine_samples:
+                ## cosine phase
+                # decay_ratio = float(self.num_steps) / float(self.lr_decay_steps)
+                # TODO
+                decay_ratio = float(self.num_steps) / float(self.stablelm2_scheduler_config.cosine_period_samples)
+                cosine_min_lr = self.stablelm2_scheduler_config.cosine_max_lr * 0.1
+                delta_lr = self.stablelm2_scheduler_config.cosine_max_lr - cosine_min_lr
+                coeff = 0.5 * (math.cos(2 * math.pi * decay_ratio) + 1.0)
+                self.stablelm2_scheduler_config.cosine_lr = cosine_min_lr + coeff * delta_lr
+                return self.stablelm2_scheduler_config.cosine_lr
+            elif self.num_steps <= self.stablelm2_scheduler_config.rsqrt_samples:
+                ## rsqrt phase
+                alpha = self.stablelm2_scheduler_config.alpha
+                beta = self.stablelm2_scheduler_config.beta
+                gbs = self.stablelm2_scheduler_config.global_batch_size * 1.0
+                self.stablelm2_scheduler_config.rsqrt_lr = alpha / ((self.num_steps / gbs + beta) ** 0.5)
+                return self.stablelm2_scheduler_config.rsqrt_lr
+            elif self.stablelm2_scheduler_config.decay_samples <= 0:
+                ## optional linear phase
+                decay_steps_ = self.lr_decay_steps - self.stablelm2_scheduler_config.rsqrt_samples
+                num_steps_ = self.num_steps - self.stablelm2_scheduler_config.rsqrt_samples
+                decay_ratio = float(num_steps_) / float(decay_steps_)
+                coeff = (1.0 - decay_ratio)
+                return coeff * self.stablelm2_scheduler_config.rsqrt_lr
+            else:
+                ## optional linear phase
+                valid_lr_decay_steps_ = min(
+                    self.lr_decay_steps,
+                    self.stablelm2_scheduler_config.rsqrt_samples + self.stablelm2_scheduler_config.decay_samples)
+                if self.num_steps <= valid_lr_decay_steps_:
+                    decay_steps_ = valid_lr_decay_steps_ - self.stablelm2_scheduler_config.rsqrt_samples
+                    num_steps_ = self.num_steps - self.stablelm2_scheduler_config.rsqrt_samples
+                    decay_ratio = float(num_steps_) / float(decay_steps_)
+                    coeff = (1.0 - decay_ratio)
+                    delta_lr = self.stablelm2_scheduler_config.rsqrt_lr - self.min_lr
+                    assert decay_ratio >= 0.0
+                    return coeff * delta_lr + self.min_lr
+                else:
+                    return self.min_lr
+
+        # Warmup-Stable-Decay(WSD)
+        if self.lr_decay_style == 'warmup-stable-decay':
+            W = self.lr_warmup_steps
+            S = round((self.lr_decay_steps - W) * 10. / 11.)
+            ## D is 10% of S.
+            T = self.lr_decay_steps - W - S
+            ## Warmup Phase, see above
+            ## Stable Phase
+            if self.num_steps < S:
+                return self.max_lr
+            else: # Decay Phase
+                return self.max_lr * 0.5 ** ((self.num_steps - S) / T)
+
         num_steps_ = self.num_steps - self.lr_warmup_steps
         decay_steps_ = self.lr_decay_steps - self.lr_warmup_steps
         decay_ratio = float(num_steps_) / float(decay_steps_)
diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py
index b9a5ef99..46a4287f 100644
--- a/megatron/core/parallel_state.py
+++ b/megatron/core/parallel_state.py
@@ -10,7 +10,9 @@ from typing import Callable, List, Optional
 
 import torch
 
-from .utils import GlobalMemoryBuffer, is_torch_min_version
+from megatron.core.utils import GlobalMemoryBuffer, is_torch_min_version
+
+from flagscale.train import get_parallel_context
 
 # Intra-layer model parallel group that the current rank belongs to.
 _TENSOR_MODEL_PARALLEL_GROUP = None
@@ -115,6 +117,8 @@ _TENSOR_AND_CONTEXT_PARALLEL_GROUP = None
 # combined parallel group of TP, DP, and CP used for fp8
 _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP = None
 
+_LAST_RANK_WHEN_USING_PIPELINE = None
+
 # Memory buffers to avoid dynamic memory allocation
 _GLOBAL_MEMORY_BUFFER = None
 
@@ -980,6 +984,8 @@ def initialize_model_parallel(
     global _POSITION_EMBEDDING_GROUP
     global _POSITION_EMBEDDING_GLOBAL_RANKS
     assert _POSITION_EMBEDDING_GROUP is None, 'position embedding group is already initialized'
+    global _LAST_RANK_WHEN_USING_PIPELINE
+    assert _LAST_RANK_WHEN_USING_PIPELINE is None, 'last rank when using pipeline 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.
@@ -1083,6 +1089,8 @@ def initialize_model_parallel(
             _POSITION_EMBEDDING_GROUP = group
             _POSITION_EMBEDDING_GLOBAL_RANKS = position_embedding_ranks
 
+    _LAST_RANK_WHEN_USING_PIPELINE = list(generator_wrapper('pp'))[-1][-1]
+
     # Build the tensor + data parallel groups.
     global _TENSOR_AND_DATA_PARALLEL_GROUP
     global _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP
@@ -1213,6 +1221,9 @@ def initialize_model_parallel(
 
 def is_initialized():
     """Useful for code segments that may be accessed with or without mpu initialization"""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return True
     return _DATA_PARALLEL_GROUP is not None
 
 
@@ -1228,6 +1239,10 @@ def is_unitialized() -> bool:
 
 def model_parallel_is_initialized():
     """Check if model- and data-parallel groups are initialized."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return True
+
     if (
         _TENSOR_MODEL_PARALLEL_GROUP is None
         or _PIPELINE_MODEL_PARALLEL_GROUP is None
@@ -1239,6 +1254,9 @@ def model_parallel_is_initialized():
 
 def get_model_parallel_group(check_initialized=True):
     """Get the model-parallel group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_model_parallel_group(check_initialized=check_initialized)
     if check_initialized:
         assert _MODEL_PARALLEL_GROUP is not None, 'model parallel group is not initialized'
     return _MODEL_PARALLEL_GROUP
@@ -1246,6 +1264,10 @@ def get_model_parallel_group(check_initialized=True):
 
 def get_tensor_model_parallel_group(check_initialized=True):
     """Get the tensor-model-parallel group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_tensor_model_parallel_group(check_initialized)
+
     if check_initialized:
         assert (
             _TENSOR_MODEL_PARALLEL_GROUP is not None
@@ -1255,6 +1277,10 @@ def get_tensor_model_parallel_group(check_initialized=True):
 
 def get_pipeline_model_parallel_group(check_initialized=True):
     """Get the pipeline-model-parallel group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_pipeline_model_parallel_group(check_initialized=check_initialized)
+
     if check_initialized:
         assert (
             _PIPELINE_MODEL_PARALLEL_GROUP is not None
@@ -1264,6 +1290,12 @@ def get_pipeline_model_parallel_group(check_initialized=True):
 
 def get_data_parallel_group(with_context_parallel=False, partial_data_parallel=False):
     """Get the data-parallel group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_data_parallel_group(
+            with_context_parallel, partial_data_parallel
+        )
+
     if with_context_parallel:
         if partial_data_parallel:
             assert (
@@ -1282,6 +1314,12 @@ def get_data_parallel_group(with_context_parallel=False, partial_data_parallel=F
 
 def get_data_parallel_group_gloo(with_context_parallel=False, partial_data_parallel=False):
     """Get the Gloo data-parallel group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_data_parallel_group_gloo(
+            with_context_parallel, partial_data_parallel
+        )
+
     if with_context_parallel:
         if partial_data_parallel:
             assert (
@@ -1300,6 +1338,10 @@ def get_data_parallel_group_gloo(with_context_parallel=False, partial_data_paral
 
 def get_inter_partial_data_parallel_group():
     """Get the group spanning the different partial data-parallel groups."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_inter_partial_data_parallel_group()
+
     assert (
         _INTER_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP is not None
     ), 'Inter partial data parallel group is not initialized'
@@ -1308,6 +1350,10 @@ def get_inter_partial_data_parallel_group():
 
 def get_context_parallel_group(check_initialized=True):
     """Get the context-parallel group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_context_parallel_group(check_initialized)
+
     if check_initialized:
         assert _CONTEXT_PARALLEL_GROUP is not None, 'context parallel group is not initialized'
     return _CONTEXT_PARALLEL_GROUP
@@ -1315,6 +1361,10 @@ def get_context_parallel_group(check_initialized=True):
 
 def get_context_parallel_global_ranks(check_initialized=True):
     """Get all global ranks of the context-parallel group that the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_context_parallel_global_ranks(check_initialized)
+
     if check_initialized:
         assert (
             _CONTEXT_PARALLEL_GLOBAL_RANKS is not None
@@ -1324,6 +1374,10 @@ def get_context_parallel_global_ranks(check_initialized=True):
 
 def get_hierarchical_context_parallel_groups(check_initialized=True):
     """Get the inner ring of context parallel group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_hierarchical_context_parallel_groups(check_initialized)
+
     if check_initialized:
         assert _HIERARCHICAL_CONTEXT_PARALLEL_GROUPS is not None
     return _HIERARCHICAL_CONTEXT_PARALLEL_GROUPS
@@ -1331,6 +1385,10 @@ def get_hierarchical_context_parallel_groups(check_initialized=True):
 
 def get_embedding_group(check_initialized=True):
     """Get the embedding group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_embedding_group(check_initialized=check_initialized)
+
     if check_initialized:
         assert _EMBEDDING_GROUP is not None, 'embedding group is not initialized'
     return _EMBEDDING_GROUP
@@ -1338,6 +1396,9 @@ def get_embedding_group(check_initialized=True):
 
 def get_position_embedding_group(check_initialized=True):
     """Get the position embedding group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_position_embedding_group(check_initialized=check_initialized)
     if check_initialized:
         assert _POSITION_EMBEDDING_GROUP is not None, 'position embedding group is not initialized'
     return _POSITION_EMBEDDING_GROUP
@@ -1345,6 +1406,10 @@ def get_position_embedding_group(check_initialized=True):
 
 def get_amax_reduction_group(with_context_parallel=False, tp_only_amax_red=False):
     """Get the FP8 amax reduction group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_amax_reduction_group(with_context_parallel)
+
     if with_context_parallel:
         if not tp_only_amax_red:
             assert (
@@ -1371,6 +1436,10 @@ def get_amax_reduction_group(with_context_parallel=False, tp_only_amax_red=False
 
 def get_tensor_and_data_parallel_group(with_context_parallel=False):
     """Get the tensor- and data-parallel group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_tensor_and_data_parallel_group(with_context_parallel)
+
     if with_context_parallel:
         assert (
             _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP is not None
@@ -1385,6 +1454,9 @@ def get_tensor_and_data_parallel_group(with_context_parallel=False):
 
 def get_tensor_and_context_parallel_group(check_initialized=True):
     """Get the tensor- and context-parallel group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_tensor_and_context_parallel_group(check_initialized=check_initialized)
     if check_initialized:
         assert (
             _TENSOR_AND_CONTEXT_PARALLEL_GROUP is not None
@@ -1394,32 +1466,52 @@ def get_tensor_and_context_parallel_group(check_initialized=True):
 
 def set_tensor_model_parallel_world_size(world_size):
     """Set the tensor-model-parallel size"""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_tensor_model_parallel_world_size(world_size)
+
     global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
     _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = world_size
 
 
 def set_pipeline_model_parallel_world_size(world_size):
     """Set the pipeline-model-parallel size"""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_pipeline_model_parallel_world_size(world_size)
+
     global _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
     _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = world_size
 
 
 def set_virtual_pipeline_model_parallel_world_size(world_size):
     """Set the pipeline-model-parallel size"""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx = para_ctx.set_virtual_pipeline_model_parallel_world_size(world_size)
+
     global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
     _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = world_size
 
 
 def get_tensor_model_parallel_world_size():
     """Return world size for the tensor-model-parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_tensor_model_parallel_world_size()
+
     global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
     if _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE is not None:
         return _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
     return torch.distributed.get_world_size(group=get_tensor_model_parallel_group())
 
 
-def get_pipeline_model_parallel_world_size():
+def get_pipeline_model_parallel_world_size(group=None):
     """Return world size for the pipeline-model-parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_pipeline_model_parallel_world_size(group)
+
     global _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
     if _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE is not None:
         return _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
@@ -1438,32 +1530,52 @@ def get_pipeline_model_parallel_world_size():
 
 def set_tensor_model_parallel_rank(rank):
     """Set tensor-model-parallel rank."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_tensor_model_parallel_rank(rank)
+
     global _MPU_TENSOR_MODEL_PARALLEL_RANK
     _MPU_TENSOR_MODEL_PARALLEL_RANK = rank
 
 
 def set_pipeline_model_parallel_rank(rank):
     """Set pipeline-model-parallel rank."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_pipeline_model_parallel_rank(rank)
+
     global _MPU_PIPELINE_MODEL_PARALLEL_RANK
     _MPU_PIPELINE_MODEL_PARALLEL_RANK = rank
 
 
 def set_pipeline_model_parallel_split_rank(rank):
     """Set pipeline-model-parallel split rank. DEPRECATED."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_pipeline_model_parallel_split_rank(rank)
+
     global _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
     _PIPELINE_MODEL_PARALLEL_SPLIT_RANK = rank
 
 
 def get_tensor_model_parallel_rank():
     """Return caller's rank for the tensor-model-parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_tensor_model_parallel_rank()
+
     global _MPU_TENSOR_MODEL_PARALLEL_RANK
     if _MPU_TENSOR_MODEL_PARALLEL_RANK is not None:
         return _MPU_TENSOR_MODEL_PARALLEL_RANK
     return torch.distributed.get_rank(group=get_tensor_model_parallel_group())
 
 
-def get_pipeline_model_parallel_rank():
+def get_pipeline_model_parallel_rank(group=None):
     """Return caller's rank for the pipeline-model-parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_pipeline_model_parallel_rank(group)
+
     global _MPU_PIPELINE_MODEL_PARALLEL_RANK
     if _MPU_PIPELINE_MODEL_PARALLEL_RANK is not None:
         return _MPU_PIPELINE_MODEL_PARALLEL_RANK
@@ -1484,12 +1596,20 @@ def get_pipeline_model_parallel_rank():
 
 def get_pipeline_model_parallel_split_rank():
     """Return pipeline-model-parallel split rank."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_pipeline_model_parallel_split_rank()
+
     global _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
     return _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
 
 
-def is_pipeline_first_stage(ignore_virtual=False):
+def is_pipeline_first_stage(ignore_virtual=False, group=None):
     """Return True if in the first pipeline model-parallel stage, False otherwise."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.is_pipeline_first_stage(ignore_virtual, group)
+
     if not ignore_virtual:
         if (
             get_virtual_pipeline_model_parallel_world_size() is not None
@@ -1499,8 +1619,12 @@ def is_pipeline_first_stage(ignore_virtual=False):
     return get_pipeline_model_parallel_rank() == 0
 
 
-def is_pipeline_last_stage(ignore_virtual=False):
+def is_pipeline_last_stage(ignore_virtual=False, group=None):
     """Return True if in the last pipeline-model-parallel stage, False otherwise."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.is_pipeline_last_stage(ignore_virtual, group)
+
     if not ignore_virtual:
         virtual_pipeline_model_parallel_world_size = (
             get_virtual_pipeline_model_parallel_world_size()
@@ -1514,8 +1638,12 @@ def is_pipeline_last_stage(ignore_virtual=False):
     return get_pipeline_model_parallel_rank() == (get_pipeline_model_parallel_world_size() - 1)
 
 
-def is_rank_in_embedding_group(ignore_virtual=False):
+def is_rank_in_embedding_group(ignore_virtual=False, group=None):
     """Return true if current rank is in embedding group, False otherwise."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.is_rank_in_embedding_group(ignore_virtual, group)
+
     rank = torch.distributed.get_rank()
     global _EMBEDDING_GLOBAL_RANKS
     if _EMBEDDING_GLOBAL_RANKS is None:
@@ -1532,16 +1660,24 @@ def is_rank_in_embedding_group(ignore_virtual=False):
     return False
 
 
-def is_rank_in_position_embedding_group():
+def is_rank_in_position_embedding_group(group=None):
     """Return true if current rank is in position embedding group, False otherwise."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.is_rank_in_position_embedding_group(group)
+
     rank = torch.distributed.get_rank()
     global _POSITION_EMBEDDING_GLOBAL_RANKS
     return _POSITION_EMBEDDING_GLOBAL_RANKS is not None and rank in _POSITION_EMBEDDING_GLOBAL_RANKS
 
 
-def is_pipeline_stage_before_split(rank=None):
+def is_pipeline_stage_before_split(rank=None, group=None):
     """Return True if pipeline stage executes encoder block for a model
     with both encoder and decoder."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.is_pipeline_stage_before_split(rank, group)
+
     if get_pipeline_model_parallel_world_size() == 1:
         return True
     if rank is None:
@@ -1554,9 +1690,13 @@ def is_pipeline_stage_before_split(rank=None):
     return False
 
 
-def is_pipeline_stage_after_split(rank=None):
+def is_pipeline_stage_after_split(rank=None, group=None):
     """Return True if pipeline stage executes decoder block for a model
     with both encoder and decoder."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.is_pipeline_stage_after_split(rank, group)
+
     if get_pipeline_model_parallel_world_size() == 1:
         return True
     if rank is None:
@@ -1573,6 +1713,10 @@ def is_inside_encoder(rank=None) -> bool:
     """Return True if pipeline stage executes encoder block.
     This function implicitly assumes we have a model with both
     encoder and decoder."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.is_inside_encoder(rank)
+
     if get_pipeline_model_parallel_world_size() == 1:
         return True
     if rank is None:
@@ -1595,6 +1739,10 @@ def is_inside_encoder(rank=None) -> bool:
 def is_inside_decoder(rank=None) -> bool:
     """Return True if pipeline stage executes decoder block for a model
     with both encoder and decoder."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.is_inside_decoder(rank)
+
     if get_pipeline_model_parallel_world_size() == 1:
         return True
     if rank is None:
@@ -1608,33 +1756,53 @@ def is_inside_decoder(rank=None) -> bool:
 
 
 def get_pipeline_model_parallel_decoder_start() -> Optional[int]:
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_pipeline_model_parallel_decoder_start()
+
     """Return decoder start rank (if encoder pipeline parallelism is set)."""
     global _PIPELINE_MODEL_PARALLEL_DECODER_START
     return _PIPELINE_MODEL_PARALLEL_DECODER_START
 
 
-def is_pipeline_stage_at_split():
+def is_pipeline_stage_at_split(group=None):
     """Return true if pipeline stage executes decoder block and next
     stage executes encoder block for a model with both encoder and
     decoder."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.is_pipeline_stage_at_split(group)
+
     rank = get_pipeline_model_parallel_rank()
     return is_pipeline_stage_before_split(rank) and is_pipeline_stage_after_split(rank + 1)
 
 
 def get_virtual_pipeline_model_parallel_rank():
     """Return the virtual pipeline-parallel rank."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_virtual_pipeline_model_parallel_rank()
+
     global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
     return _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
 
 
 def set_virtual_pipeline_model_parallel_rank(rank):
     """Set the virtual pipeline-parallel rank."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_virtual_pipeline_model_parallel_rank(rank)
+
     global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
     _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = rank
 
 
 def get_virtual_pipeline_model_parallel_world_size():
     """Return the virtual pipeline-parallel world size."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_virtual_pipeline_model_parallel_world_size()
+
     global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
     return _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
 
@@ -1642,6 +1810,10 @@ def get_virtual_pipeline_model_parallel_world_size():
 def get_tensor_model_parallel_src_rank():
     """Calculate the global rank corresponding to the first local rank
     in the tensor model parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_tensor_model_parallel_src_rank()
+
     assert (
         _TENSOR_MODEL_PARALLEL_GLOBAL_RANKS is not None
     ), "Tensor model parallel group is not initialized"
@@ -1651,6 +1823,10 @@ def get_tensor_model_parallel_src_rank():
 def get_model_parallel_src_rank():
     """Calculate the global rank corresponding to the first local rank
     in the model parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_model_parallel_src_rank()
+
     assert _MODEL_PARALLEL_GLOBAL_RANKS is not None, "Model parallel group is not initialized"
     return _MODEL_PARALLEL_GLOBAL_RANKS[0]
 
@@ -1658,6 +1834,10 @@ def get_model_parallel_src_rank():
 def get_data_parallel_src_rank(with_context_parallel=False):
     """Calculate the global rank corresponding to the first local rank
     in the data parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_data_parallel_src_rank(with_context_parallel)
+
     if with_context_parallel:
         assert (
             _DATA_PARALLEL_GLOBAL_RANKS_WITH_CP is not None
@@ -1668,8 +1848,12 @@ def get_data_parallel_src_rank(with_context_parallel=False):
         return _DATA_PARALLEL_GLOBAL_RANKS[0]
 
 
-def get_pipeline_model_parallel_first_rank():
+def get_pipeline_model_parallel_first_rank(group=None):
     """Return the global rank of the first stage in the current rank's pipeline."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_pipeline_model_parallel_first_rank(group)
+
     assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
     if isinstance(_PIPELINE_GLOBAL_RANKS[0], list):
         # I assume the first rank is the same for all pp groups right now.
@@ -1680,8 +1864,12 @@ def get_pipeline_model_parallel_first_rank():
         return _PIPELINE_GLOBAL_RANKS[0]
 
 
-def get_pipeline_model_parallel_last_rank():
+def get_pipeline_model_parallel_last_rank(group=None):
     """Return the global rank of the last stage in the current rank's pipeline."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_pipeline_model_parallel_last_rank(group)
+
     assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
     last_rank_local = get_pipeline_model_parallel_world_size() - 1
     if isinstance(_PIPELINE_GLOBAL_RANKS[0], list):
@@ -1690,12 +1878,16 @@ def get_pipeline_model_parallel_last_rank():
         return _PIPELINE_GLOBAL_RANKS[last_rank_local]
 
 
-def get_pipeline_model_parallel_next_rank():
+def get_pipeline_model_parallel_next_rank(group=None):
     """Return the global rank that follows the caller in the pipeline, for each
     pipeline-parallel group that the rank is part of.
 
     If it is just part of one group, an int is returned, otherwise a list of ints.
     """
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_pipeline_model_parallel_next_rank(group)
+
     assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
     rank_in_pipeline = get_pipeline_model_parallel_rank()
     world_size = get_pipeline_model_parallel_world_size()
@@ -1708,12 +1900,16 @@ def get_pipeline_model_parallel_next_rank():
         return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline + 1) % world_size]
 
 
-def get_pipeline_model_parallel_prev_rank():
+def get_pipeline_model_parallel_prev_rank(group=None):
     """Return the global rank that precedes the caller in the pipeline, for each
     pipeline-parallel group that the rank is part of.
 
     If it is just part of one group, an int is returned, otherwise a list of ints.
     """
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_pipeline_model_parallel_prev_rank(group)
+
     assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
     rank_in_pipeline = get_pipeline_model_parallel_rank()
     world_size = get_pipeline_model_parallel_world_size()
@@ -1726,8 +1922,24 @@ def get_pipeline_model_parallel_prev_rank():
         return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline - 1) % world_size]
 
 
+def get_last_rank_when_using_pipeline():
+    """Return the global rank of the last process in the pipeline"""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_last_rank_when_using_pipeline()
+
+    assert _LAST_RANK_WHEN_USING_PIPELINE is not None, "Last rank when using pipeline is not initialized"
+    return _LAST_RANK_WHEN_USING_PIPELINE
+
+
 def get_data_parallel_world_size(with_context_parallel=False, partial_data_parallel=False):
     """Return world size for the data parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_data_parallel_world_size(
+            with_context_parallel, partial_data_parallel
+        )
+
     global _MPU_DATA_PARALLEL_WORLD_SIZE
     if _MPU_DATA_PARALLEL_WORLD_SIZE is not None:
         return _MPU_DATA_PARALLEL_WORLD_SIZE
@@ -1744,12 +1956,22 @@ def get_data_parallel_world_size(with_context_parallel=False, partial_data_paral
 
 def set_data_parallel_rank(rank):
     """Return world size for the data parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_data_parallel_rank(rank)
+
     global _MPU_DATA_PARALLEL_RANK
     _MPU_DATA_PARALLEL_RANK = rank
 
 
 def get_data_parallel_rank(with_context_parallel=False, partial_data_parallel=False):
     """Return caller's rank in the data-parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_data_parallel_rank(
+            with_context_parallel, partial_data_parallel
+        )
+
     global _MPU_DATA_PARALLEL_RANK
     if _MPU_DATA_PARALLEL_RANK is not None:
         return _MPU_DATA_PARALLEL_RANK
@@ -1766,6 +1988,10 @@ def get_data_parallel_rank(with_context_parallel=False, partial_data_parallel=Fa
 
 def get_context_parallel_world_size():
     """Return world size for the context parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_context_parallel_world_size()
+
     if torch.distributed.is_available() and torch.distributed.is_initialized():
         return torch.distributed.get_world_size(group=get_context_parallel_group())
     else:
@@ -1774,6 +2000,10 @@ def get_context_parallel_world_size():
 
 def get_context_parallel_rank():
     """Return caller's rank in the context-parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_context_parallel_rank()
+
     if torch.distributed.is_available() and torch.distributed.is_initialized():
         return torch.distributed.get_rank(group=get_context_parallel_group())
     else:
@@ -1782,6 +2012,10 @@ def get_context_parallel_rank():
 
 def get_tensor_and_context_parallel_world_size():
     """Return world size for the tensor and context-parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_tensor_and_context_parallel_world_size()
+
     if torch.distributed.is_available() and torch.distributed.is_initialized():
         return torch.distributed.get_world_size(group=get_tensor_and_context_parallel_group())
     else:
@@ -1790,6 +2024,10 @@ def get_tensor_and_context_parallel_world_size():
 
 def get_tensor_and_context_parallel_rank():
     """Return caller's rank in the joint tensor-model-parallel and context-parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_tensor_and_context_parallel_world_size()
+
     if torch.distributed.is_available() and torch.distributed.is_initialized():
         return torch.distributed.get_rank(group=get_tensor_and_context_parallel_group())
     else:
@@ -1799,6 +2037,10 @@ def get_tensor_and_context_parallel_rank():
 ### Expert-related parallel states functions
 def get_expert_model_parallel_group(check_initialized=True):
     """Get the expert-model-parallel group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_model_parallel_group(check_initialized)
+
     if check_initialized:
         assert (
             _EXPERT_MODEL_PARALLEL_GROUP is not None
@@ -1807,6 +2049,10 @@ def get_expert_model_parallel_group(check_initialized=True):
 
 
 def get_expert_model_parallel_world_size():
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_model_parallel_world_size()
+
     """Return world size for the expert-model-parallel group."""
     if _MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE is not None:
         return _MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE
@@ -1818,12 +2064,20 @@ def get_expert_model_parallel_world_size():
 
 def set_expert_model_parallel_world_size(world_size):
     """Sets the expert-model-parallel world size."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_expert_model_parallel_world_size(world_size)
+
     global _MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE
     _MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE = world_size
 
 
 def get_expert_model_parallel_rank():
     """Return caller's rank in the expert-model-parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_model_parallel_rank()
+
     if _MPU_EXPERT_MODEL_PARALLEL_RANK is not None:
         return _MPU_EXPERT_MODEL_PARALLEL_RANK
     if torch.distributed.is_available() and torch.distributed.is_initialized():
@@ -1834,12 +2088,20 @@ def get_expert_model_parallel_rank():
 
 def set_expert_model_parallel_rank(rank):
     """Set expert-model-parallel rank."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_expert_model_parallel_rank(rank)
+
     global _MPU_EXPERT_MODEL_PARALLEL_RANK
     _MPU_EXPERT_MODEL_PARALLEL_RANK = rank
 
 
 def get_expert_tensor_parallel_group(check_initialized=True):
     """Get the expert-tensor-parallel group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_tensor_parallel_group(check_initialized)
+
     if check_initialized:
         assert (
             _EXPERT_TENSOR_PARALLEL_GROUP is not None
@@ -1849,6 +2111,10 @@ def get_expert_tensor_parallel_group(check_initialized=True):
 
 def get_expert_tensor_parallel_world_size():
     """Return world size for the expert tensor parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_tensor_parallel_world_size()
+
     global _MPU_EXPERT_TENSOR_PARALLEL_WORLD_SIZE
     if _MPU_EXPERT_TENSOR_PARALLEL_WORLD_SIZE is not None:
         return _MPU_EXPERT_TENSOR_PARALLEL_WORLD_SIZE
@@ -1861,12 +2127,20 @@ def get_expert_tensor_parallel_world_size():
 
 def set_expert_tensor_parallel_world_size(world_size):
     "Set expert tensor model parallel size"
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_expert_tensor_parallel_world_size(world_size)
+
     global _MPU_EXPERT_TENSOR_PARALLEL_WORLD_SIZE
     _MPU_EXPERT_TENSOR_PARALLEL_WORLD_SIZE = world_size
 
 
 def get_expert_tensor_parallel_rank():
     """Return my rank for the expert tensor parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_tensor_parallel_rank()
+
     global _MPU_EXPERT_TENSOR_PARALLEL_RANK
     if _MPU_EXPERT_TENSOR_PARALLEL_RANK is not None:
         return _MPU_EXPERT_TENSOR_PARALLEL_RANK
@@ -1879,12 +2153,20 @@ def get_expert_tensor_parallel_rank():
 
 def set_expert_tensor_parallel_rank(rank):
     "Set expert tensor model parallel rank"
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_expert_tensor_parallel_rank(rank)
+
     global _MPU_EXPERT_TENSOR_PARALLEL_RANK
     _MPU_EXPERT_TENSOR_PARALLEL_RANK = rank
 
 
 def get_expert_tensor_and_model_parallel_group(check_initialized=True):
     """Get the expert-tensor and expert-model group the caller rank belongs to."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_tensor_and_model_parallel_group(check_initialized)
+
     if check_initialized:
         assert (
             _EXPERT_TENSOR_AND_MODEL_PARALLEL_GROUP is not None
@@ -1894,6 +2176,10 @@ def get_expert_tensor_and_model_parallel_group(check_initialized=True):
 
 def get_expert_tensor_and_model_parallel_world_size():
     """Return world size for the expert model parallel group times expert tensor parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_tensor_and_model_parallel_world_size()
+
     if torch.distributed.is_available() and torch.distributed.is_initialized():
         world_size = torch.distributed.get_world_size(
             group=get_expert_tensor_and_model_parallel_group()
@@ -1905,6 +2191,10 @@ def get_expert_tensor_and_model_parallel_world_size():
 
 def get_expert_tensor_and_model_parallel_rank():
     """Return caller's rank in the joint tensor- and expert-model-parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_tensor_and_model_parallel_rank()
+
     if torch.distributed.is_available() and torch.distributed.is_initialized():
         return torch.distributed.get_rank(group=get_expert_tensor_and_model_parallel_group())
     else:
@@ -1913,6 +2203,10 @@ def get_expert_tensor_and_model_parallel_rank():
 
 def get_expert_tensor_model_pipeline_parallel_group(check_initialized=True):
     """Get expert tensor-model-pipeline parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_tensor_model_pipeline_parallel_group(check_initialized=check_initialized)
+
     if check_initialized:
         assert (
             _EXPERT_TENSOR_MODEL_PIPELINE_PARALLEL_GROUP is not None
@@ -1922,6 +2216,10 @@ def get_expert_tensor_model_pipeline_parallel_group(check_initialized=True):
 
 def get_expert_data_parallel_group():
     """Get expert data parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_data_parallel_group()
+
     assert _EXPERT_DATA_PARALLEL_GROUP is not None, 'Expert data parallel group is not initialized'
     return _EXPERT_DATA_PARALLEL_GROUP
 
@@ -1933,11 +2231,19 @@ def get_data_modulo_expert_parallel_group():
         "get_expert_data_parallel_group instead.",
         DeprecationWarning,
     )
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_data_parallel_group()
+
     return get_expert_data_parallel_group()
 
 
 def get_expert_data_parallel_group_gloo():
     """Get expert data parallel group-gloo."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_data_parallel_group_gloo()
+
     assert (
         _EXPERT_DATA_PARALLEL_GROUP_GLOO is not None
     ), 'Expert data parallel group-gloo is not initialized'
@@ -1946,6 +2252,10 @@ def get_expert_data_parallel_group_gloo():
 
 def get_expert_data_parallel_rank():
     """Return caller's rank in the expert data parallel group."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_expert_data_parallel_rank()
+
     if torch.distributed.is_available() and torch.distributed.is_initialized():
         return torch.distributed.get_rank(group=get_expert_data_parallel_group())
     else:
@@ -1965,6 +2275,10 @@ def get_expert_data_parallel_world_size():
 
 def _set_global_memory_buffer():
     """Initialize global buffer."""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.set_global_memory_buffer()
+
     global _GLOBAL_MEMORY_BUFFER
     assert _GLOBAL_MEMORY_BUFFER is None, 'global memory buffer is already initialized'
     _GLOBAL_MEMORY_BUFFER = GlobalMemoryBuffer()
@@ -1972,12 +2286,19 @@ def _set_global_memory_buffer():
 
 def get_global_memory_buffer():
     """Return the global GlobalMemoryBuffer object"""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        return para_ctx.get_global_memory_buffer()
+
     assert _GLOBAL_MEMORY_BUFFER is not None, 'global memory buffer is not initialized'
     return _GLOBAL_MEMORY_BUFFER
 
 
 def destroy_global_memory_buffer():
     """Sets the global memory buffer to None"""
+    para_ctx = get_parallel_context()
+    if para_ctx is not None:
+        para_ctx.destroy_global_memory_buffer()
     global _GLOBAL_MEMORY_BUFFER
     _GLOBAL_MEMORY_BUFFER = None
 
@@ -2119,3 +2440,6 @@ def destroy_model_parallel():
         torch.distributed.destroy_process_group(_EXPERT_DATA_PARALLEL_GROUP_GLOO)
     _EXPERT_DATA_PARALLEL_GROUP_GLOO = None
     # End of expert parallelism destroy.
+
+    global _LAST_RANK_WHEN_USING_PIPELINE
+    _LAST_RANK_WHEN_USING_PIPELINE = None
diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py
index 26a96f45..e485af60 100644
--- a/megatron/core/pipeline_parallel/p2p_communication.py
+++ b/megatron/core/pipeline_parallel/p2p_communication.py
@@ -242,6 +242,7 @@ def _communicate(
     tensor_shape: Shape,
     config: ModelParallelConfig,
     wait_on_reqs: bool = True,
+    group: torch.distributed.ProcessGroup = None
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Communicate tensors between stages. Used as helper method in other
     communication methods that are used in megatron/schedules.py.
@@ -291,7 +292,7 @@ def _communicate(
         return torch.empty(
             recv_prev_shape,
             requires_grad=True,
-            device=torch.cuda.current_device(),
+            device=torch.cuda.current_device() if "cpu:gloo" != group.name() else torch.device("cpu"),
             dtype=config.pipeline_dtype,
         )
 
@@ -299,7 +300,7 @@ def _communicate(
         return torch.empty(
             recv_next_shape,
             requires_grad=True,
-            device=torch.cuda.current_device(),
+            device=torch.cuda.current_device() if "cpu:gloo" != group.name() else torch.device("cpu"),
             dtype=config.pipeline_dtype,
         )
 
@@ -342,9 +343,14 @@ def _communicate(
     # tensor parallelism, and hence a rank in the encoder is going to feed
     # several different decoder ranks. We therefore have to receive or send tensors
     # from several groups. For convenience, I wrap everything into lists.
-    pp_group = get_pipeline_model_parallel_group()
-    next_rank = get_pipeline_model_parallel_next_rank()
-    prev_rank = get_pipeline_model_parallel_prev_rank()
+    if config.enable_hetero: # Using the passed 'group' in the case of 'enable_hetero'
+        pp_group = group
+        next_rank = get_pipeline_model_parallel_next_rank(group=group)
+        prev_rank = get_pipeline_model_parallel_prev_rank(group=group)
+    else:
+        pp_group = get_pipeline_model_parallel_group()
+        next_rank = get_pipeline_model_parallel_next_rank()
+        prev_rank = get_pipeline_model_parallel_prev_rank()
     if not isinstance(pp_group, list):
         pp_group = [pp_group]
         assert not isinstance(next_rank, list)
@@ -425,12 +431,19 @@ def _communicate(
 
     return tensor_recv_prev, tensor_recv_next, reqs
 
+def warm_up_comm_group(config):
+    if config.enable_hetero:
+        from flagscale.train.hetero.p2p_communication import warm_up_comm_group_hetero
+        warm_up_comm_group_hetero(config)
 
 def recv_forward(tensor_shape: Shape, config: ModelParallelConfig) -> torch.Tensor:
     """Receive tensor from previous rank in pipeline (forward receive).
 
     See _communicate for argument details.
     """
+    if config.enable_hetero:
+        from flagscale.train.hetero.p2p_communication import recv_forward_hetero
+        return recv_forward_hetero(tensor_shape, config)
 
     if core.parallel_state.is_pipeline_first_stage():
         input_tensor = None
@@ -455,6 +468,10 @@ def recv_backward(tensor_shape: Shape, config: ModelParallelConfig) -> torch.Ten
 
     See _communicate for argument details.
     """
+    if config.enable_hetero:
+        from flagscale.train.hetero.p2p_communication import recv_backward_hetero
+        return recv_backward_hetero(tensor_shape, config)
+
     if core.parallel_state.is_pipeline_last_stage():
         output_tensor_grad = None
     else:
@@ -478,6 +495,10 @@ def send_forward(output_tensor: torch.Tensor, config: ModelParallelConfig) -> No
 
     See _communicate for argument details.
     """
+    if config.enable_hetero:
+        from flagscale.train.hetero.p2p_communication import send_forward_hetero
+        send_forward_hetero(output_tensor, config)
+        return
 
     if not core.parallel_state.is_pipeline_last_stage():
         if config.timers is not None:
@@ -499,6 +520,11 @@ def send_backward(input_tensor_grad: torch.Tensor, config: ModelParallelConfig)
 
     See _communicate for argument details.
     """
+    if config.enable_hetero:
+        from flagscale.train.hetero.p2p_communication import send_backward_hetero
+        send_backward_hetero(input_tensor_grad, config)
+        return
+
     if not core.parallel_state.is_pipeline_first_stage():
         if config.timers is not None:
             config.timers('backward-send', log_level=2).start()
@@ -521,6 +547,10 @@ def send_forward_recv_backward(
 
     See _communicate for argument details.
     """
+    if config.enable_hetero:
+        from flagscale.train.hetero.p2p_communication import send_forward_recv_backward_hetero
+        return send_forward_recv_backward_hetero(output_tensor, tensor_shape, config)
+
     if core.parallel_state.is_pipeline_last_stage():
         output_tensor_grad = None
     else:
@@ -546,6 +576,10 @@ def send_backward_recv_forward(
 
     See _communicate for argument details.
     """
+    if config.enable_hetero:
+        from flagscale.train.hetero.p2p_communication import send_backward_recv_forward_hetero
+        return send_backward_recv_forward_hetero(input_tensor_grad, tensor_shape, config)
+
     if core.parallel_state.is_pipeline_first_stage():
         input_tensor = None
     else:
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index 0e15e399..ed33de78 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -1851,6 +1851,7 @@ def forward_backward_pipelining_without_interleaving(
         output_tensors = []
     forward_data_store = []
 
+    p2p_communication.warm_up_comm_group(config=config)
     # Run warmup forward passes.
     for i in range(num_warmup_microbatches):
         # Decide to checkpoint all layers' activations of the current micro-batch
diff --git a/megatron/core/timers.py b/megatron/core/timers.py
index 27195169..f4dfa063 100644
--- a/megatron/core/timers.py
+++ b/megatron/core/timers.py
@@ -280,6 +280,8 @@ class Timers:
                 # groups inside their class.
                 rank_name_to_time[rank, i] = self._timers[name].elapsed(reset=reset)
 
+        if "cpu:gloo" == torch.distributed.get_backend():
+            rank_name_to_time = rank_name_to_time.cpu()
         # See the note above for why we are not using gather.
         dist_all_gather_func(rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1))
 
diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
index c087d5d0..88ac379b 100644
--- a/megatron/core/transformer/attention.py
+++ b/megatron/core/transformer/attention.py
@@ -722,22 +722,44 @@ class SelfAttention(Attention):
         )
 
         if submodules.q_layernorm is not None:
-            self.q_layernorm = build_module(
-                submodules.q_layernorm,
-                hidden_size=self.hidden_size_per_attention_head,
-                config=self.config,
-                eps=self.config.layernorm_epsilon,
-            )
+            if not self.config.qk_layernorm_hidden_dim:
+                self.q_layernorm = build_module(
+                    submodules.q_layernorm,
+                    hidden_size=self.hidden_size_per_attention_head,
+                    config=self.config,
+                    eps=self.config.layernorm_epsilon,
+                )
+            else:
+                tp_world_size = get_tensor_model_parallel_world_size()
+                assert tp_world_size <= 1, "TP world size must be less than 1 for qk_layernorm_hidden_dim"
+                # nums_head_cur_rank = divide(self.config.num_attention_heads, tp_world_size)
+                self.q_layernorm = build_module(
+                    submodules.q_layernorm,
+                    hidden_size=self.query_projection_size,
+                    config=self.config,
+                    eps=self.config.layernorm_epsilon,
+                )
         else:
             self.q_layernorm = None
 
         if submodules.k_layernorm is not None:
-            self.k_layernorm = build_module(
-                submodules.k_layernorm,
-                hidden_size=self.hidden_size_per_attention_head,
-                config=self.config,
-                eps=self.config.layernorm_epsilon,
-            )
+            if not self.config.qk_layernorm_hidden_dim:
+                self.k_layernorm = build_module(
+                    submodules.k_layernorm,
+                    hidden_size=self.hidden_size_per_attention_head,
+                    config=self.config,
+                    eps=self.config.layernorm_epsilon,
+                )
+            else:
+                tp_world_size = get_tensor_model_parallel_world_size()
+                assert tp_world_size <= 1, "TP world size must be less than 1 for qk_layernorm_hidden_dim"
+                # nums_head_cur_rank = divide(self.config.num_attention_heads, tp_world_size)
+                self.k_layernorm = build_module(
+                    submodules.k_layernorm,
+                    hidden_size=self.kv_projection_size,
+                    config=self.config,
+                    eps=self.config.layernorm_epsilon,
+                )
         else:
             self.k_layernorm = None
 
@@ -854,10 +876,24 @@ class SelfAttention(Attention):
         query = query.reshape(query.size(0), query.size(1), -1, self.hidden_size_per_attention_head)
 
         if self.q_layernorm is not None:
-            query = self.q_layernorm(query)
+            if not self.config.qk_layernorm_hidden_dim:
+                query = self.q_layernorm(query)
+            else:
+                # [sq, b, np, hn] -> [sq, b, 1, np * hn]
+                query_shape = list(query.shape)
+                query = query.reshape(query.size(0), query.size(1), 1, -1)
+                query = self.q_layernorm(query)
+                query = query.reshape(*query_shape)
 
         if self.k_layernorm is not None:
-            key = self.k_layernorm(key)
+            if not self.config.qk_layernorm_hidden_dim:
+                key = self.k_layernorm(key)
+            else:
+                # [sq, b, ng, hn] -> [sq, b, 1, ng * hn]
+                key_shape = list(key.shape)
+                key = key.reshape(key.size(0), key.size(1), 1, -1)
+                key = self.k_layernorm(key)
+                key = key.reshape(*key_shape)
 
         if self.config.test_mode:
             self.run_realtime_tests()
diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py
index 0f0cf863..44bb4dd0 100644
--- a/megatron/core/transformer/moe/experts.py
+++ b/megatron/core/transformer/moe/experts.py
@@ -225,6 +225,12 @@ class GroupedMLP(MegatronModule):
         setattr(self.weight1, 'allreduce', not self.expert_parallel)
         setattr(self.weight2, 'allreduce', not self.expert_parallel)
 
+        # NOTE(lizhiyu): The following code is for hetro-expert training when one of the expert parallel degree is 1.
+        #                 But there are other codes need to be modified to make it work.
+        # if config.enable_hetero:
+        #     setattr(self.weight1, 'allreduce', False)
+        #     setattr(self.weight2, 'allreduce', False)
+
         def remove_extra_states_check(self, incompatible_keys):
             """
             Remove _extra_state from unexpected keys.
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 31dec135..bed4d436 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -657,6 +657,39 @@ def reduce_aux_losses_tracker_across_ranks(track_names: Optional[List[str]] = No
             )
 
 
+def reduce_aux_losses_tracker_across_ranks_hetero(
+    track_names: Optional[List[str]] = None,
+):
+    """Collect and reduce the auxiliary losses across ranks."""
+    tracker = parallel_state.get_moe_layer_wise_logging_tracker()
+    if track_names is None:
+        track_names = tracker.keys()
+    for name in track_names:
+        values = tracker[name]["values"]
+        # Reduce aux losses across ranks.
+        if tracker[name].get("reduce_group") is not None:
+            torch.distributed.all_reduce(
+                values, group=tracker[name].get("reduce_group")
+            )
+        if tracker[name].get("avg_group") is not None:
+            torch.distributed.all_reduce(
+                values,
+                group=tracker[name]["avg_group"],
+                op=torch.distributed.ReduceOp.AVG,
+            )
+        pp_groups = parallel_state.get_pipeline_model_parallel_group()
+        if "cpu:gloo" == torch.distributed.get_backend(pp_groups[0]):
+            values = values.cpu()
+        assert isinstance(pp_groups, list), "pp_groups should be a list for hetero."
+        if len(pp_groups) > 1:
+            origin_values = values.clone().detach()
+            for pp_group in pp_groups:
+                values.copy_(origin_values)
+                torch.distributed.all_reduce(values, group=pp_group)
+        else:
+            torch.distributed.all_reduce(values, group=pp_groups[0])
+
+
 def track_moe_metrics(
     loss_scale: float,
     iteration: int,
@@ -668,6 +701,7 @@ def track_moe_metrics(
     track_names: Optional[List[str]] = None,
     num_layers: Optional[int] = None,
     moe_layer_freq: Optional[Union[int, List[int]]] = None,
+    enable_hetero=False,
 ):
     """Track the MoE metrics for logging."""
     # Aux loss logging
@@ -681,7 +715,10 @@ def track_moe_metrics(
                     tracker[key]["values"] = torch.zeros(num_layers, device="cuda")
                     tracker[key]["reduce_group"] = None
                     tracker[key]["avg_group"] = None
-    reduce_aux_losses_tracker_across_ranks(track_names)
+    if not enable_hetero:
+        reduce_aux_losses_tracker_across_ranks(track_names)
+    else:
+        reduce_aux_losses_tracker_across_ranks_hetero(track_names)
 
     # Get number of MoE layers
     if moe_layer_freq is None:
diff --git a/megatron/core/transformer/torch_norm.py b/megatron/core/transformer/torch_norm.py
index d0ceca7a..72504cfa 100644
--- a/megatron/core/transformer/torch_norm.py
+++ b/megatron/core/transformer/torch_norm.py
@@ -69,7 +69,7 @@ class L2Norm(torch.nn.Module):
         self.hidden_size = hidden_size
         self.eps = eps
 
-    @jit_fuser
+    #@jit_fuser
     def _norm(self, x):
         """
         Performs the actual L2 normalization.
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
old mode 100755
new mode 100644
index 38a400a4..7edb62e1
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -1,5 +1,5 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-
+import os
 from contextlib import nullcontext
 from dataclasses import dataclass
 from typing import List, Optional, Union
@@ -103,7 +103,7 @@ def get_num_layers_to_build(config: TransformerConfig) -> int:
             and config.num_layers_in_last_pipeline_stage is not None
         ):
             num_layers_per_pipeline_rank = config.num_layers_in_last_pipeline_stage
-    else:
+    elif not config.enable_hetero:
         # Include the embedding layer and loss layer into pipeline parallelism partition
         num_layers = config.num_layers
         if config.account_for_embedding_in_pipeline_split:
@@ -120,7 +120,7 @@ def get_num_layers_to_build(config: TransformerConfig) -> int:
     if (
         parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None
         and config.pipeline_model_parallel_size > 1
-    ):
+    ) and parallel_state.get_virtual_pipeline_model_parallel_world_size() > 1:
         # Interleaved pipeline parallelism:
         # Number of layers in each model chunk is the number of layers in the stage,
         # divided by the number of model chunks in a stage.
@@ -145,7 +145,12 @@ def get_num_layers_to_build(config: TransformerConfig) -> int:
     else:
         # Non-interleaved pipeline parallelism:
         # Each stage gets a contiguous set of layers.
-        num_layers_to_build = num_layers_per_pipeline_rank
+
+        if config.enable_hetero:
+            pipeline_rank = parallel_state.get_pipeline_model_parallel_rank()
+            num_layers_to_build = config.hetero_pipeline_layer_split[pipeline_rank]
+        else:
+            num_layers_to_build = num_layers_per_pipeline_rank
 
     # The embedding (or loss) layer cannot function as a standalone transformer layer
     # Reduce the number of layers to construct by 1 on the first (or last) stage if the
@@ -472,6 +477,14 @@ class TransformerBlock(MegatronModule):
             [s, b, h], and optionally the updated context tensor if cross-attention is used.
         """
 
+        ########## FlagScale Begin ##########
+        # for refined recompute
+        if hasattr(self.layers[0], 'current_microbatch'):
+            self.current_microbatch = self.layers[0].current_microbatch
+        else: # multimodal model
+            self.current_microbatch = -1
+        ########## FlagScale End ##########
+
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
         # Delete the obsolete reference to the initial input tensor if necessary
@@ -517,6 +530,71 @@ class TransformerBlock(MegatronModule):
         use_inner_fp8_context = self.config.fp8 and self.config.fp8_recipe != Fp8Recipe.delayed
         outer_fp8_context = get_fp8_context(self.config) if use_outer_fp8_context else nullcontext()
 
+        if self.config.recompute_method_per_stage_micro_batch != None:
+            if self.config.virtual_pipeline_model_parallel_size != None:
+                if (
+                    self.config.recompute_method_per_stage_micro_batch[
+                        parallel_state.get_virtual_pipeline_model_parallel_rank()
+                        * self.config.pipeline_model_parallel_size
+                        + parallel_state.get_pipeline_model_parallel_rank()
+                    ][self.current_microbatch]
+                    == 0
+                ):
+                    self.config.recompute_method = 'uniform'
+                elif (
+                    self.config.recompute_method_per_stage_micor_batch[
+                        parallel_state.get_virtual_pipeline_model_parallel_rank()
+                        * self.config.pipeline_model_parallel_size
+                        + parallel_state.get_pipeline_model_parallel_rank()
+                    ][self.current_microbatch]
+                    == 1
+                ):
+                    self.config.recompute_method = 'block'
+                else:
+                    raise ValueError("the item of recompute_method_per_stage_micor_batch must be '0' or '1' ")
+            else:
+                if (
+                    self.config.recompute_method_per_stage_micro_batch[
+                        parallel_state.get_pipeline_model_parallel_rank()
+                    ][self.current_microbatch]
+                    == 0
+                ):
+                    self.config.recompute_method = 'uniform'
+                elif (
+                    self.config.recompute_method_per_stage_micro_batch[
+                        parallel_state.get_pipeline_model_parallel_rank()
+                    ][self.current_microbatch]
+                    == 1
+                ):
+                    self.config.recompute_method = 'block'
+                else:
+                    raise ValueError("the item of recompute_method_per_stage_micor_batch must be '0' or '1' ")
+
+        if self.config.recompute_num_layers_per_stage_micro_batch != None:
+            if self.config.virtual_pipeline_model_parallel_size != None:
+                self.config.recompute_num_layers = self.config.recompute_num_layers_per_stage_micro_batch[
+                    parallel_state.get_virtual_pipeline_model_parallel_rank()
+                    * self.config.pipeline_model_parallel_size
+                    + parallel_state.get_pipeline_model_parallel_rank()
+                ][self.current_microbatch]
+            else:
+                self.config.recompute_num_layers = self.config.recompute_num_layers_per_stage_micro_batch[
+                    parallel_state.get_pipeline_model_parallel_rank()
+                ][self.current_microbatch]
+            if self.config.recompute_num_layers == 0:
+                self.config.recompute_method = None
+                self.config.recompute_granularity = None
+
+        if (
+            self.config.recompute_granularity_per_stage_micro_batch != None
+            and self.config.recompute_granularity_per_stage_micro_batch[
+                parallel_state.get_pipeline_model_parallel_rank()
+            ][self.current_microbatch]
+            == 0
+        ):
+            self.config.recompute_granularity = None
+            self.config.recompute_method = None
+
         with rng_context, outer_fp8_context:
             # Forward pass.
             if self.config.recompute_granularity == 'full' and self.training:
@@ -596,6 +674,16 @@ class TransformerBlock(MegatronModule):
         elif isinstance(self.config.moe_layer_freq, list):
             non_homogeneous_layers = True
 
+
+        # TODO: @aoyulong - This is a temporary solution to support single-file-per-tensor ckpt
+        non_homogeneous_layers_env = os.getenv('FS_NON_HOMOGENEOUS_LAYERS', 'False').lower() in (
+            'true',
+            '1',
+            't',
+        )
+        if non_homogeneous_layers_env:
+            non_homogeneous_layers = True
+
         if self.config.heterogeneous_block_specs:
             non_homogeneous_layers = True
 
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index c4bfe416..94da18d3 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -232,6 +232,15 @@ class TransformerConfig(ModelParallelConfig):
     the number of transformer layers to recompute within each pipeline stage.  Must be None for
     'selective' activation checkpointing."""
 
+    recompute_granularity_per_stage_micro_batch: list = None
+    """Same as recompute_granularity but for each stage and each micro-batch."""
+
+    recompute_method_per_stage_micro_batch: list = None
+    """Same as recompute_method but for each stage and each micro-batch."""
+
+    recompute_num_layers_per_stage_micro_batch: list = None
+    """Same as recompute_num_layers but for each stage and each micro-batch."""
+
     distribute_saved_activations: Optional[bool] = None
     """If True, distribute recomputed activations across the model parallel group."""
 
@@ -509,6 +518,12 @@ class TransformerConfig(ModelParallelConfig):
     config_logger_dir: str = ""
     """When non-empty, dumps entry-point configs to config_logger_dir"""
 
+    qk_layernorm_hidden_dim: bool = False
+    """Whether to apply LayerNorm to the query and key embeddings on the hidden dimension rather than head dimension."""
+
+    use_partial_reduce_for_shared_embedding: bool = False
+    """Whether to use partional reduce for shared embedding."""
+
     flash_decode: bool = False
     """ Use the optimized flash decoding kernel during inference. """
 
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index b7d761f8..ca7de92c 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -31,8 +31,9 @@ def get_transformer_layer_offset(config: TransformerConfig):
             pipeline_rank = pipeline_rank - pp_decoder_start
 
     if config.pipeline_model_parallel_size > 1:
-
-        if (
+        if config.enable_hetero:
+            offset = sum(([0] + config.hetero_pipeline_layer_split)[: pipeline_rank + 1])
+        elif (
             config.num_layers_in_first_pipeline_stage is not None
             or config.num_layers_in_last_pipeline_stage is not None
         ):
diff --git a/megatron/inference/text_generation/sampling.py b/megatron/inference/text_generation/sampling.py
index 370773a3..b2fc7573 100644
--- a/megatron/inference/text_generation/sampling.py
+++ b/megatron/inference/text_generation/sampling.py
@@ -42,7 +42,7 @@ def modify_logits_for_top_p_filtering(logits, top_p):
 
 
 
-def sample(logits, top_k=0, top_p=0.0, temperature=1.0, vocab_size=None):
+def sample(logits, top_k=0, top_p=0.0, temperature=1.0, vocab_size=None, seed=123):
     """ Sample and generate a token.
     Note: logits has the dimension [b, v] where b is the batch size
           and v is the vocabulary size.
@@ -51,6 +51,9 @@ def sample(logits, top_k=0, top_p=0.0, temperature=1.0, vocab_size=None):
     generations due to padding.
     """
 
+    generator = torch.Generator(device=logits.device)
+    generator.manual_seed(seed)
+
     # Check logits for consistency.
     assert logits.ndim == 2, 'expected the logits to be of [b, v] shape.'
     assert logits.type() == 'torch.cuda.FloatTensor', \
@@ -83,7 +86,7 @@ def sample(logits, top_k=0, top_p=0.0, temperature=1.0, vocab_size=None):
 
         # After filtering, we need to recalculate the distribution.
         probs = logits.softmax(dim=-1)
-        samples = torch.multinomial(probs, num_samples=1).view(-1)
+        samples = torch.multinomial(probs, num_samples=1, generator=generator).view(-1)
 
     # If vocab size is provided, make sure the samples are in
     # in the range [0, vocab-size).
diff --git a/megatron/inference/text_generation/tokenization.py b/megatron/inference/text_generation/tokenization.py
index 541cc47b..a41c7f59 100644
--- a/megatron/inference/text_generation/tokenization.py
+++ b/megatron/inference/text_generation/tokenization.py
@@ -41,6 +41,15 @@ def detokenize_generations(tokens_gpu_tensor,
                     word = bytearray([tokenizer.tokenizer.byte_decoder[c] for c in word]).decode(
                         "utf-8", errors="replace"
                     )
+                    args = get_args()
+                    if args.tokenizer_type == 'AquilaTokenizer':
+                        if token in tokenizer.tokenizer.special_tokens_decoder:
+                            word = tokenizer.tokenizer.special_tokens_decoder[token]
+                        else :
+                            word = tokenizer.tokenizer.decoder[token]
+                            word = bytearray(
+                                [tokenizer.tokenizer.byte_decoder[c] for c in word]).decode(
+                                    'utf-8', errors='replace')
                     words.append(word)
 
             prompts_plus_generations_segments.append(words)
diff --git a/megatron/legacy/fused_kernels/__init__.py b/megatron/legacy/fused_kernels/__init__.py
index 87cceac3..5a04def1 100644
--- a/megatron/legacy/fused_kernels/__init__.py
+++ b/megatron/legacy/fused_kernels/__init__.py
@@ -56,7 +56,7 @@ def load(args):
 
 def _get_cuda_bare_metal_version(cuda_dir):
     raw_output = subprocess.check_output(
-        [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
+        [cuda_dir + "/bin/cucc", "-V"], universal_newlines=True
     )
     output = raw_output.split()
     release_idx = output.index("release") + 1
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index c5d750e8..33926361 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -53,6 +53,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser):
     parser = _add_autoresume_args(parser)
     parser = _add_biencoder_args(parser)
     parser = _add_vision_args(parser)
+    parser = _add_mtp_args(parser)
     parser = _add_moe_args(parser)
     parser = _add_mla_args(parser)
     parser = _add_heterogeneous_args(parser)
@@ -77,6 +78,9 @@ def parse_args(extra_args_provider=None, ignore_unknown_args=False):
                                      allow_abbrev=False)
 
     parser = add_megatron_arguments(parser)
+    parser = _add_hetero_args(parser)
+    parser = _add_auto_tuner_args(parser)
+    parser = _add_auto_skip_spiky_loss(parser)
 
     # Custom arguments.
     if extra_args_provider is not None:
@@ -308,66 +312,74 @@ def validate_args(args, defaults={}):
             "legacy model format only supports the 'torch' checkpoint format."
     update_use_dist_ckpt(args)
 
-    if args.encoder_pipeline_model_parallel_size == 0 and args.num_experts == 0:
-        assert args.encoder_tensor_model_parallel_size == args.tensor_model_parallel_size,  "If non-MOE encoder shares first decoder pipeline rank it must have the same TP as the decoder."
-
-    if args.encoder_tensor_model_parallel_size > 0:
-        assert args.num_attention_heads % args.encoder_tensor_model_parallel_size == 0
-        assert args.encoder_tensor_model_parallel_size <= args.tensor_model_parallel_size, "We do not support encoders with more TP than the decoder."
+    if args.attention_backend == AttnBackend.local:
+            assert args.spec[0] == 'local' , '--attention-backend local is only supported with --spec local'
 
-    if args.encoder_pipeline_model_parallel_size > 0 and args.encoder_tensor_model_parallel_size == 0:
-        args.encoder_tensor_model_parallel_size = args.tensor_model_parallel_size
+    if not args.enable_hetero:
+        if args.encoder_pipeline_model_parallel_size == 0 and args.num_experts == 0:
+            assert args.encoder_tensor_model_parallel_size == args.tensor_model_parallel_size,  "If non-MOE encoder shares first decoder pipeline rank it must have the same TP as the decoder."
 
-    encoder_model_size = args.encoder_tensor_model_parallel_size * args.encoder_pipeline_model_parallel_size * args.context_parallel_size
-    decoder_model_size = args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size
-    total_model_size = encoder_model_size + decoder_model_size
+        if args.encoder_tensor_model_parallel_size > 0:
+            assert args.num_attention_heads % args.encoder_tensor_model_parallel_size == 0
+            assert args.encoder_tensor_model_parallel_size <= args.tensor_model_parallel_size, "We do not support encoders with more TP than the decoder."
 
-    # Total model size.
-    assert args.world_size % total_model_size == 0, (
-        f"world size ({args.world_size}) is not divisible by total_model_size ({encoder_model_size=} + {decoder_model_size=})"
-    )
-
-    if args.attention_backend == AttnBackend.local:
-        assert args.spec[0] == 'local' , '--attention-backend local is only supported with --spec local'
+        if args.encoder_pipeline_model_parallel_size > 0 and args.encoder_tensor_model_parallel_size == 0:
+            args.encoder_tensor_model_parallel_size = args.tensor_model_parallel_size
 
-    # Pipeline model parallel size.
-    args.transformer_pipeline_model_parallel_size = args.pipeline_model_parallel_size
+        encoder_model_size = args.encoder_tensor_model_parallel_size * args.encoder_pipeline_model_parallel_size * args.context_parallel_size
+        decoder_model_size = args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size
+        total_model_size = encoder_model_size + decoder_model_size
 
-    args.data_parallel_size = args.world_size // total_model_size
-
-    if args.rank == 0:
-        print('using world size: {}, data-parallel size: {}, '
-              'context-parallel size: {}, '
-              'hierarchical context-parallel sizes: {}'
-              'tensor-model-parallel size: {}, '
-              'encoder-tensor-model-parallel size: {}, '
-              'pipeline-model-parallel size: {}, '
-              'encoder-pipeline-model-parallel size: {}'.format(
-                  args.world_size, args.data_parallel_size,
-                  args.context_parallel_size,
-                  args.hierarchical_context_parallel_sizes,
-                  args.tensor_model_parallel_size,
-                  args.encoder_tensor_model_parallel_size,
-                  args.pipeline_model_parallel_size,
-                  args.encoder_pipeline_model_parallel_size), flush=True)
+        # Total model size.
+        assert args.world_size % total_model_size == 0, (
+            f"world size ({args.world_size}) is not divisible by total_model_size ({encoder_model_size=} + {decoder_model_size=})"
+        )
 
-    # Checks.
+        if args.attention_backend == AttnBackend.local:
+            assert args.spec[0] == 'local' , '--attention-backend local is only supported with --spec local'
 
-    # Backwards compatibility.
-    if args.pipeline_model_parallel_split_rank is not None:
-        args.encoder_pipeline_model_parallel_size = args.pipeline_model_parallel_split_rank
-        args.pipeline_model_parallel_size -= args.encoder_pipeline_model_parallel_size
-        assert args.pipeline_model_parallel_size > 0
+        # Pipeline model parallel size.
+        args.transformer_pipeline_model_parallel_size = (
+            args.pipeline_model_parallel_size - 1
+            if args.standalone_embedding_stage else
+            args.pipeline_model_parallel_size
+        )
 
-    if args.hierarchical_context_parallel_sizes:
-        from numpy import prod
-        assert args.context_parallel_size == prod(args.hierarchical_context_parallel_sizes)
-    if "a2a+p2p" in args.cp_comm_type:
-        assert args.hierarchical_context_parallel_sizes is not None, \
-        "--hierarchical-context-parallel-sizes must be set when a2a+p2p is used in cp comm"
+        args.data_parallel_size = args.world_size // total_model_size
 
-    if args.expert_tensor_parallel_size is None:
-        args.expert_tensor_parallel_size = args.tensor_model_parallel_size
+        if args.rank == 0:
+            print('using world size: {}, data-parallel size: {}, '
+                  'context-parallel size: {}, '
+                  'hierarchical context-parallel sizes: {}'
+                  'tensor-model-parallel size: {}, '
+                  'encoder-tensor-model-parallel size: {}, '
+                  'pipeline-model-parallel size: {}, '
+                  'encoder-pipeline-model-parallel size: {}'.format(
+                      args.world_size, args.data_parallel_size,
+                      args.context_parallel_size,
+                      args.hierarchical_context_parallel_sizes,
+                      args.tensor_model_parallel_size,
+                      args.encoder_tensor_model_parallel_size,
+                      args.pipeline_model_parallel_size,
+                      args.encoder_pipeline_model_parallel_size), flush=True)
+
+        # Checks.
+
+        # Backwards compatibility.
+        if args.pipeline_model_parallel_split_rank is not None:
+            args.encoder_pipeline_model_parallel_size = args.pipeline_model_parallel_split_rank
+            args.pipeline_model_parallel_size -= args.encoder_pipeline_model_parallel_size
+            assert args.pipeline_model_parallel_size > 0
+
+        if args.hierarchical_context_parallel_sizes:
+            from numpy import prod
+            assert args.context_parallel_size == prod(args.hierarchical_context_parallel_sizes)
+        if "a2a+p2p" in args.cp_comm_type:
+            assert args.hierarchical_context_parallel_sizes is not None, \
+            "--hierarchical-context-parallel-sizes must be set when a2a+p2p is used in cp comm"
+
+        if args.expert_tensor_parallel_size is None:
+            args.expert_tensor_parallel_size = args.tensor_model_parallel_size
 
     # Deprecated arguments.
     assert args.batch_size is None, '--batch-size argument is no longer ' \
@@ -441,6 +453,7 @@ def validate_args(args, defaults={}):
         '--num-layers-per-virtual-pipeline-stage and --num-virtual-stages-per-pipeline-rank cannot be set at the same time'
 
     if args.num_layers_per_virtual_pipeline_stage is not None or args.num_virtual_stages_per_pipeline_rank is not None:
+        assert args.enable_hetero is False, 'num_layers_per_virtual_pipeline_stage is not supported with heterogeneous parallelism for now'
         if args.overlap_p2p_comm:
             assert args.pipeline_model_parallel_size > 1, \
                 'When interleaved schedule is used, pipeline-model-parallel size '\
@@ -498,8 +511,9 @@ def validate_args(args, defaults={}):
                 if args.account_for_loss_in_pipeline_split:
                     num_layers += 1
 
-                assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \
-                    'Number of layers should be divisible by the pipeline-model-parallel size'
+                if args.enable_hetero is False:
+                    assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \
+                        'Number of layers should be divisible by the pipeline-model-parallel size'
     if args.rank == 0:
         print(f"Number of virtual stages per pipeline stage: {args.virtual_pipeline_model_parallel_size}")
 
@@ -606,6 +620,11 @@ def validate_args(args, defaults={}):
             if args.rank == 0:
                 print('accumulate and all-reduce gradients in fp32 for '
                       'bfloat16 data type.', flush=True)
+        if args.optimizer == 'hybridadam':
+            args.accumulate_bf16 = False
+        if args.accumulate_bf16:
+            args.accumulate_allreduce_grads_in_fp32 = False
+
 
     if args.rank == 0:
         print('using {} for parameters ...'.format(args.params_dtype),
@@ -678,12 +697,22 @@ def validate_args(args, defaults={}):
     # Checks.
     if args.ffn_hidden_size is None:
         if args.swiglu:
-            # reduce the dimnesion for MLP since projections happens on
-            # two linear layers. this keeps the number of paramters in
-            # the same ballpark as the counterpart with 4*h size
-            # we keep it a multiple of 64, which means the actual tensor size
-            # will be a multiple of 64 / tp_size
-            args.ffn_hidden_size = int((4 * args.hidden_size * 2 / 3) / 64) * 64
+            # Ref: https://github.com/facebookresearch/llama/blob/main/llama/model.py#L161-L162
+            if args.multiple_of is not None:
+                hidden_dim = int(4 * args.hidden_size * 2 / 3)
+                if args.hidden_dim_multiplier is not None:
+                    assert args.hidden_dim_multiplier > 0, \
+                        'multiplier for hidden dim should be greater than zero'
+                    hidden_dim = int(hidden_dim * args.hidden_dim_multiplier)
+                args.ffn_hidden_size = args.multiple_of * \
+                    ((hidden_dim + args.multiple_of - 1) // args.multiple_of)
+            else:
+                # reduce the dimnesion for MLP since projections happens on
+                # two linear layers. this keeps the number of paramters in
+                # the same ballpark as the counterpart with 4*h size
+                # we keep it a multiple of 64, which means the actual tensor size
+                # will be a multiple of 64 / tp_size
+                args.ffn_hidden_size = int((4 * args.hidden_size * 2 / 3) / 64) * 64
         else:
             args.ffn_hidden_size = 4 * args.hidden_size
 
@@ -1302,6 +1331,8 @@ def _add_network_size_args(parser):
                        help='Which normalization technique to use.')
     group.add_argument('--norm-epsilon', type=float, default=1e-5,
                        help='Epsilon for layer norm and RMS norm.')
+    group.add_argument('--norm-init-weight', type=float, default=None,
+                       help="Norm weight initialization.")
     group.add_argument('--apply-layernorm-1p', action='store_true',
                        help='Adjust LayerNorm weights such that they are centered '
                        'around zero. This improves numerical stability.')
@@ -1317,6 +1348,10 @@ def _add_network_size_args(parser):
                        help='Use squared relu activation instead of default gelu')
     group.add_argument('--swiglu', action='store_true',
                        help='Use gated linear units and SiLU activation instead of default gelu')
+    group.add_argument('--multiple-of', type=int, default=None,
+                       help='Multiplier for setting Feed-Forward Network hidden size when swiglu.')
+    group.add_argument('--hidden-dim-multiplier', type=float, default=None,
+                       help='Custom Multiplier for setting Feed-Forward Network hidden dim when swiglu.')
     group.add_argument('--onnx-safe', type=bool, required=False,
                        help='Use workarounds for known problems with '
                        'Torch ONNX exporter')
@@ -1473,8 +1508,18 @@ def _add_logging_args(parser):
                        help='The wandb experiment name.')
     group.add_argument('--wandb-save-dir', type=str, default='',
                        help='Path to save the wandb results locally.')
+    group.add_argument('--wandb-mode', type=str, choices=['online', 'offline', 'disabled'], default='offline',
+                       help='Can be "online", "offline" or "disabled". Defaults to "offline".')
+    group.add_argument('--wandb-api-key', type=str, default='',
+                       help='The wandb API keys and must be provided if using online mode.')
+    group.add_argument('--wandb-log-model', action='store_true',
+                       help='If set, write model to wandb.')
+    group.add_argument('--wandb-log-model-interval', type=int, default=1000,
+                       help='The interval to save the model to wandb.')
     group.add_argument('--logging-level', type=int, default=None,
                        help='Set default logging level')
+    group.add_argument('--device-tflops', type=int, default=0,
+                       help='tflops for device, if set will print mfu')
     return parser
 
 
@@ -1596,6 +1641,25 @@ def _add_training_args(parser):
                        '"moe": recompute the MoE layer.'
                        '"moe_act", "layernorm", and "mla_up_proj" use output-discarding checkpointing, '
                        '"core_attn", "mlp", and "moe" uses normal checkpointing.')
+    group.add_argument('--recompute-granularity-per-stage-micro-batch', nargs='*', type=str, default=None,
+                       help='used with recompute-granularity=full, setting recompute granularity'
+                       'of each stage and each micro-batch. This argument must be a two-dimension list, '
+                       'the sum of the first item of all the sub-lists should be equal to pipeline-model-parallel-size.'
+                       'Every sub-list is in the form: n0, flag0, n1, flag1,... except the first item, which is the stage number.'
+                       'The sum of n0, n1, ... should be equal to nums-micro-batch.'
+                       'granularity flag: 0 means turning off full recompute, 1 means turning on')
+    group.add_argument('--recompute-method-per-stage-micro-batch', nargs='*', type=str, default=None,
+                       help='used with recompute-granularity=full, setting recompute method '
+                       'of each stage and each micro-batch. This argument must be a two-dimension list, '
+                       'the sum of the first item of all the sub-lists should be equal to pipeline-model-parallel-size.'
+                       'Every sub-list is in the form: n0, flag0, n1, flag1,... except the first item, which is the stage number.'
+                       'The sum of n0, n1, ... should be equal to nums-micro-batch.'
+                       'method: 0 means uniform, 1 means block')
+    group.add_argument('--recompute-num-layers-per-stage-micro-batch', nargs='*', type=str, default=None,
+                       help='used with recompute-granularity=full, setting recompute num layers '
+                       'of each stage and each micro-batch. This argument must be a two-dimension list, '
+                       'Every sub-list is in the form: n0, num_laryers0, n1, num_laryers1,... except the first item, which is the stage number.'
+                       'The sum of n0, n1, ... should be equal to nums-micro-batch. ')
     group.add_argument('--no-clone-scatter-output-in-embedding', action='store_false',
                        help='If not set, clone the output of the scatter in embedding layer to GC original tensor.',
                        dest='clone_scatter_output_in_embedding')
@@ -1682,6 +1746,10 @@ def _add_training_args(parser):
                        help='Total number of samples to train over all '
                        'training runs. Note that either train-iters or '
                        'train-samples should be provided.')
+    group.add_argument('--skip-samples-range', nargs='+', type=int, default=None,
+                       help='Range of samples to skip during training.')
+    group.add_argument('--skip-iters-range', nargs='+', type=int, default=None,
+                       help='Range of iterations to skip during training.')
     group.add_argument('--log-interval', type=int, default=100,
                        help='Report loss and timing interval.')
     group.add_argument('--exit-interval', type=int, default=None,
@@ -1795,6 +1863,8 @@ def _add_training_args(parser):
                        choices=['nccl', 'ucc'],
                        help='Select a communicator backend for pipeline parallel communication. '
                        'If None, the default backend will be used.')
+    group.add_argument('--accumulate-bf16', action='store_true',
+                       help='use bf16 accumulate force')
 
     return parser
 
@@ -1843,11 +1913,26 @@ def _add_learning_rate_args(parser):
                        'and initial warmup, the learning rate at each '
                        'iteration would be different.')
     group.add_argument('--lr-decay-style', type=str, default='linear',
-                       choices=['constant', 'linear', 'cosine', 'inverse-square-root', 'WSD'],
+                       choices=['constant', 'linear', 'cosine', 'inverse-square-root', 'WSD', 'stablelm2-scheduler'],
                        help='Learning rate decay function.')
     group.add_argument('--lr-wsd-decay-style', type=str, default='exponential',
                        choices=['exponential', 'linear', 'cosine'],
                        help='Decay style for the annealing phase of WSD'),
+    ## stablelm2-scheduler consists of multiple stages
+    group.add_argument('--lr-decay-stablelm2-cosine-samples', type=int, default=0,
+                       help='Samples number of cosine scheduler including warmup samples, used in stablelm2 scheduler.')
+    group.add_argument('--lr-decay-stablelm2-cosine-max-lr', type=float, default=None,
+                       help='Maximum lr of cosine scheduler, used in stablelm2 scheduler.')
+    group.add_argument('--lr-decay-stablelm2-cosine-period-samples', type=int, default=0,
+                       help='Period of cosine scheduler, used in stablelm2 scheduler.')
+    group.add_argument('--lr-decay-stablelm2-rsqrt-samples', type=int, default=0,
+                       help='Samples number of rsqrt scheduler used in stablelm2 scheduler.')
+    group.add_argument('--lr-decay-stablelm2-decay-samples', type=int, default=0,
+                       help='Samples number of decay scheduler used in stablelm2 scheduler.')
+    group.add_argument('--lr-decay-stablelm2-alpha', type=float, default=1.0,
+                       help='Numerator used in stablelm2 scheduler.')
+    group.add_argument('--lr-decay-stablelm2-beta', type=float, default=0.0,
+                       help='Denominator used in stablelm2 scheduler.')
     group.add_argument('--lr-decay-iters', type=int, default=None,
                        help='number of iterations to decay learning rate over,'
                        ' If None defaults to `--train-iters`')
@@ -1903,6 +1988,8 @@ def _add_checkpointing_args(parser):
                        help='Output directory to save checkpoints to.')
     group.add_argument('--save-interval', '--persistent-save-interval', type=int, default=None,
                        help='Number of iterations between persistent checkpoint saves.')
+    group.add_argument('--rampup-save-interval', type=int, default=None,
+                       help='Number of iterations between checkpoint saves.in the ramup phase.')
     group.add_argument('--no-save-optim', action='store_true', default=None,
                        help='Do not save current optimizer.')
     group.add_argument('--no-save-rng', action='store_true', default=None,
@@ -1948,6 +2035,8 @@ def _add_checkpointing_args(parser):
     group.add_argument('--no-use-tokenizer-model-from-checkpoint-args', action='store_false',
                        dest='use_tokenizer_model_from_checkpoint_args',
                        help='If set, do not use tokenizer model path from checkpoint')
+    group.add_argument('--save-when-num-microbatches-change', action='store_true',
+                       help='Save param name to index maps only')
     group.add_argument('--exit-on-missing-checkpoint', action='store_true',
                        help="If '--load' is set, but checkpoint is not found "
                        "(e.g., path typo), then exit instead of random "
@@ -2090,7 +2179,7 @@ def _add_distributed_args(parser):
                        default=False, help='if set, overlap pipeline parallel communication in warmup and flush',
                        dest='overlap_p2p_comm_warmup_flush')
     group.add_argument('--distributed-backend', default='nccl',
-                       choices=['nccl', 'gloo'],
+                       choices=['nccl', 'gloo', 'flagcx'],
                        help='Which backend to use for distributed training.')
     group.add_argument('--distributed-timeout-minutes', type=int, default=10,
                        help='Timeout minutes for torch.distributed.')
@@ -2141,6 +2230,11 @@ def _add_distributed_args(parser):
                        'complete it instead. Also turns on '
                        '--use-cpu-initialization flag. This is for '
                        'external DDP manager.' )
+    group.add_argument('--standalone-embedding-stage', action='store_true',
+                       default=False, help='If set, *input* embedding layer '
+                       'is placed on its own pipeline stage, without any '
+                       'transformer layers. (For T5, this flag currently only '
+                       'affects the encoder embedding.)')
     group.add_argument('--account-for-embedding-in-pipeline-split', action='store_true',
                        default=False, help='If set, *input* embedding layer will be treated as a standard transformer'
                        'layer in the context of partition and placement for pipeline parallelism.')
@@ -2164,6 +2258,10 @@ def _add_distributed_args(parser):
                         'and performance requirements.')
     group.add_argument('--keep-fp8-transpose-cache-when-using-custom-fsdp', action='store_true',
                        help='If set, keep the fp8 transpose cache when using custom FSDP.')
+    group.add_argument('--use-partial-reduce-for-shared-embedding', action='store_true',
+                       help='Use partial reduce for shared word embedding.')
+    group.add_argument('--no-shared-fs', action='store_true',
+                       help='Indicate whether not running on a shared file system.')
     group.add_argument('--num-distributed-optimizer-instances', type=int, default=1,
                        help='Number of Distributed Optimizer copies across Data Parallel domain.')
     group.add_argument('--use-torch-fsdp2', action='store_true',
@@ -2213,6 +2311,9 @@ def _add_validation_args(parser):
     group.add_argument('--eval-interval', type=int, default=1000,
                        help='Interval between running evaluation on '
                        'validation set.')
+    group.add_argument('--extra-eval-interval', type=int, default=None,
+                       help='Interval between running evaluation on '
+                       'extra validation sets.')
     group.add_argument("--test-mode", action="store_true", help='Run all real-time test alongside the experiment.')
     group.add_argument('--skip-train', action='store_true',
                        default=False, help='If set, bypass the training loop, '
@@ -2227,6 +2328,8 @@ def _add_tokenizer_args(parser):
                        help='Size of vocab before EOD or padding.')
     group.add_argument('--vocab-file', type=str, default=None,
                        help='Path to the vocab file.')
+    group.add_argument('--special-tokens-file', type=str, default=None,
+                       help='Path to the BPE special tokens file.')
     group.add_argument('--merge-file', type=str, default=None,
                        help='Path to the BPE merge file.')
     group.add_argument('--vocab-extra-ids', type=int, default=0,
@@ -2244,8 +2347,17 @@ def _add_tokenizer_args(parser):
                                 'TikTokenizer',
                                 'MultimodalTokenizer',
                                 'NullTokenizer',
-                                'NullMultimodalTokenizer'],
+                                'NullMultimodalTokenizer',
+                                'AquilaTokenizerFS',
+                                'HFTokenizerFS',
+                                'HFTokenizersTokenizerFS',
+                                'Llama3TokenizerFS',
+                                'QwenTokenizerFS',
+                                'Qwen2TokenizerFS',
+                                'Qwen2VLTokenizer',],
                        help='What type of tokenizer to use.')
+    group.add_argument('--tokenizer-path', type=str, default=None,
+                       help='Path to the huggingface tokenizer.')
     group.add_argument('--tokenizer-model', type=str, default=None,
                        help='Sentencepiece tokenizer model.')
     group.add_argument('--tiktoken-pattern', type=str, default=None,
@@ -2279,6 +2391,11 @@ def _add_data_args(parser):
     group.add_argument('--valid-data-path', nargs='*', default=None,
                        help='The weight and prefix list for an independent validation dataset. '
                        'Follows the same pattern rules as --data-path.')
+    group.add_argument('--extra-valid-data-path', nargs='*', default=None,
+                       help='The weight, prefix list for an independent extra validation dataset. '
+                       'The accepted format is a list of weight, prefix and tag, '
+                       'e.g. weight1 prefix1 tag1 weight2 prefix2 tag2. '
+                       'The weight1 means the number of tokens in the prefix1 dataset. ')
     group.add_argument('--test-data-path', nargs='*', default=None,
                        help='The weight and prefix list for an independent test dataset. '
                        'Follows the same pattern rules as --data-path.')
@@ -2327,11 +2444,18 @@ def _add_data_args(parser):
                        'end-of-document token.')
     group.add_argument('--eod-mask-loss', action='store_true',
                        help='Mask loss for the end of document tokens.')
+    group.add_argument('--finetune-dataset-type', type=str, default=None,
+                       choices=['CPT', None],
+                       help='datasets type during finetunning.')
     group.add_argument('--no-create-attention-mask-in-dataloader', action='store_false',
                        help='If set, do not create attention_masks in dataloader.',
                        dest='create_attention_mask_in_dataloader')
     group.add_argument('--num-dataset-builder-threads', type=int, default=1,
                        help='Number of parallel threads per rank for dataset builder')
+    group.add_argument('--apply-sft-dataset-separated-loss-mask-if-existed', action='store_true',
+                       help='If set, use sft dataset with separated loss mask files, '
+                       'if _loss_mask_document.bin and _loss_mask_document.idx existed.')
+
     group.add_argument('--object-storage-cache-path', type=str, default=None,
                        help='Path to cache index files when using s3 or msc dataloader')
     group.add_argument('--mid-level-dataset-surplus', type=float, default=0.005,
@@ -2408,6 +2532,19 @@ def _add_biencoder_args(parser):
     return parser
 
 
+def _add_mtp_args(parser):
+    # add args for Multi-token Prediction module
+    group = parser.add_argument_group(title="mtp")
+
+    # general mtp arguements
+    group.add_argument('--num-mtp-predictor', type=int, default=0,
+                       help='num of multi token predictors')
+    group.add_argument('--mtp-loss-coeff', type=float, default=0.3,
+                       help='Scaling coefficient for mtp loss: 0.3 is recommended in DeepSeekV3.')
+
+    return parser
+
+
 def _add_vision_args(parser):
     group = parser.add_argument_group(title="vision")
 
@@ -2476,6 +2613,8 @@ def _add_vision_args(parser):
     # regularization arguments
     group.add_argument('--qk-layernorm', action='store_true',
                        help='Whether to layer normalize the q and k attention embeddings.')
+    group.add_argument('--qk-layernorm-hidden-dim', action='store_true',
+                       help='Whether to layer normalize the q and k attention embeddings on hidden dimension rather than head dimension')
 
     return parser
 
@@ -2714,4 +2853,45 @@ def _add_msc_args(parser):
     group = parser.add_argument_group(title="msc")
     group.add_argument('--disable-msc', default=True, action='store_false', dest='enable_msc',
                        help='Disable the usage of Multi-Storage Client (MSC) in Megatron Core.')
-    return parser
\ No newline at end of file
+    return parser
+
+
+def _add_hetero_args(parser):
+    group = parser.add_argument_group(title="heterogeneous training")
+
+    group.add_argument('--enable-hetero', action="store_true",
+                       help='the mode of heterogeneous training')
+    group.add_argument('--hetero-device-types', nargs='*', type=str, default=None,
+                       help='the list of device types: device_type_0 device_type_1 ...')
+    group.add_argument('--hetero-current-device-type', type=str, default=None,
+                       help='the current device type')
+    group.add_argument('--hetero-pipeline-layer-split', nargs='*', type=int, default=None,
+                       help='Incompatible with --num-layers-per-virtual-pipeline-stage for now.'
+                       'hetero-pipeline-layer-split must be in the form: layers_0 layers_1 ... layers_n. The number of the list should be equal to pipeline-model-parallel-size.')
+    group.add_argument('--hetero-process-meshes', nargs='*', type=int, default=None,
+                       help='Use this arg to set TP-CP-DP-PP of each process mesh.'
+                       'This argument must be in the form: TP0, CP0, DP0, PP0, TP1, CP0, DP1, PP1...TPN, CPN, DPN, PPN. CP and TP size can be different, sum of PP should match pipeline-model-parallel-size, DP size should be the same.')
+    group.add_argument('--expert-tensor-parallel-size-per-process-mesh', nargs='*', type=int, default=None,
+                       help='The number of tensor parallel experts for each process-mesh. The number of the list should be equal to the number of process-meshes.')
+    group.add_argument('--hetero-use-cpu-communication', action='store_true', help='Use CPU for communication for heterogeneous communication.')
+
+    return parser
+
+
+def _add_auto_tuner_args(parser):
+    group = parser.add_argument_group(title="auto tuner")
+
+    group.add_argument('--auto-tune', action='store_true',
+                       help='use auto tuner')
+
+    return parser
+
+
+def _add_auto_skip_spiky_loss(parser):
+    group = parser.add_argument_group(title='auto skip spiky loss')
+
+    group.add_argument('--auto-skip-spiky-loss', action='store_true',
+                       help='Automatically skip spiky loss iterations.')
+    group.add_argument('--spiky-loss-threshold', type=float, default=0.2,
+                          help='Threshold for skipping spiky loss iterations.')
+    return parser
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index a706c181..c1a27231 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -239,12 +239,14 @@ def read_metadata(tracker_filename):
                 print_rank_0('ERROR: Invalid metadata file {}. Exiting'.format(
                     tracker_filename))
                 sys.exit()
-    assert iteration > 0 or release, 'error parsing metadata file {}'.format(
+    # TODO: we use iteration 0 to load checkpoint from other framework.
+    # We should remove this after we have a better way to load checkpoint from other framework.
+    assert iteration >= 0 or release, 'error parsing metadata file {}'.format(
         tracker_filename)
 
     # Get the max iteration retrieved across the ranks.
     if torch.distributed.is_initialized():
-        iters_cuda = torch.tensor([iteration], dtype=torch.long, device='cuda')
+        iters_cuda = torch.tensor([iteration], dtype=torch.long, device='cuda' if 'nccl' in torch.distributed.get_backend() else 'cpu')
         torch.distributed.all_reduce(iters_cuda, op=torch.distributed.ReduceOp.MAX)
         max_iter = iters_cuda[0].item()
 
@@ -1056,6 +1058,10 @@ def load_args_from_checkpoint(
             checkpoint_args, 'add_bias_linear', not getattr(checkpoint_args, 'disable_bias_linear')
         )
 
+    # For backward compatibility.
+    if hasattr(checkpoint_args, 'apply_layernorm_rms'):
+        checkpoint_args.normalization = 'RMSNorm'
+
     def _set_arg(arg_name, old_arg_name=None, force=False):
         if not force and getattr(args, arg_name, None) is not None:
             return
@@ -1091,6 +1097,8 @@ def load_args_from_checkpoint(
     _set_arg('add_qkv_bias', force=True)
     _set_arg('squared_relu', force=True)
     _set_arg('swiglu', force=True)
+    _set_arg('multiple_of', force=True)
+    _set_arg('hidden_dim_multiplier', force=True)
     _set_arg('untie_embeddings_and_output_weights', force=True)
     _set_arg('apply_layernorm_1p', force=True)
     _set_arg('normalization', force=True)
diff --git a/megatron/training/global_vars.py b/megatron/training/global_vars.py
index 70701341..d8b3ba17 100644
--- a/megatron/training/global_vars.py
+++ b/megatron/training/global_vars.py
@@ -5,12 +5,15 @@
 import os
 import sys
 import torch
+import torch.distributed
 
 from megatron.core import Timers
 from megatron.core.num_microbatches_calculator import init_num_microbatches_calculator, unset_num_microbatches_calculator
 from megatron.training import dist_signal_handler
 from megatron.training.tokenizer import build_tokenizer
 
+from flagscale.train.hetero.p2p_communication import get_device_type_for_comm
+
 _GLOBAL_ARGS = None
 _GLOBAL_TOKENIZER = None
 _GLOBAL_TENSORBOARD_WRITER = None
@@ -91,9 +94,6 @@ def set_global_variables(args, build_tokenizer=True):
     )
     if build_tokenizer:
         _ = _build_tokenizer(args)
-    _set_tensorboard_writer(args)
-    _set_wandb_writer(args)
-    _set_one_logger(args)
     _set_adlr_autoresume(args)
     _set_timers(args)
 
@@ -101,6 +101,42 @@ def set_global_variables(args, build_tokenizer=True):
         _set_signal_handler()
 
 
+def set_global_writers(args):
+    """Set tensorboard-writer and wandb writer.
+
+    Note that this function should be called after calling finish_mpu_init.
+    This is because we can know which rank is the last one after the rank mapping in finish_mpu_init.
+    """
+
+    assert args is not None
+
+    _ensure_var_is_initialized(_GLOBAL_ARGS, 'args')
+
+    from .utils import is_last_rank
+    if is_last_rank():
+        _set_tensorboard_writer(args)
+        _set_one_logger(args)
+
+    # build wandb writers for all processes in the dp group of the last rank
+    from megatron.core import mpu
+    mp_groups = mpu.get_model_parallel_group()
+    if not isinstance(mp_groups, list):
+        mp_groups = [mp_groups]
+    size = torch.distributed.get_world_size(mp_groups[-1])
+    comm_device = get_device_type_for_comm(mp_groups)
+    ranks_tensor = torch.tensor([0 for _ in range(size)], dtype=torch.int, device=comm_device)
+    orig_ranks = torch.tensor([i for i in range(size)], dtype=torch.int, device=comm_device)
+    if is_last_rank():
+        ranks_list = torch.distributed.get_process_group_ranks(mp_groups[-1])
+        ranks_tensor = torch.tensor(ranks_list, dtype=torch.int, device=comm_device)
+    orig_ranks = ranks_tensor.clone().detach()
+    for group in mp_groups:
+        ranks_tensor = orig_ranks.clone()
+        torch.distributed.all_reduce(ranks_tensor, group=group)
+    if torch.distributed.get_rank() in ranks_tensor.tolist():
+        _set_wandb_writer(args)
+
+
 def unset_global_variables():
     """Unset global vars.
 
@@ -156,7 +192,7 @@ def _set_tensorboard_writer(args):
                                    'tensorboard writer')
 
     if hasattr(args, 'tensorboard_dir') and \
-       args.tensorboard_dir and args.rank == (args.world_size - 1):
+       args.tensorboard_dir:
         try:
             from torch.utils.tensorboard import SummaryWriter
             print('> setting tensorboard ...')
@@ -173,22 +209,37 @@ def _set_wandb_writer(args):
     global _GLOBAL_WANDB_WRITER
     _ensure_var_is_not_initialized(_GLOBAL_WANDB_WRITER,
                                    'wandb writer')
-    if getattr(args, 'wandb_project', '') and args.rank == (args.world_size - 1):
+    if getattr(args, 'wandb_project', ''):
         if args.wandb_exp_name == '':
             raise ValueError("Please specify the wandb experiment name!")
 
         import wandb
+        rank = torch.distributed.get_rank()
+
         if args.wandb_save_dir:
             save_dir = args.wandb_save_dir
         else:
             # Defaults to the save dir.
             save_dir = os.path.join(args.save, 'wandb')
+        save_dir = os.path.join(save_dir, "rank-{}".format(rank))
+        os.makedirs(save_dir, exist_ok=True)
+
+        wandb_id = f"{args.wandb_exp_name}-rank-{rank}"
+        name = f'{args.wandb_exp_name}-rank-{rank}'
+        group = args.wandb_exp_name
         wandb_kwargs = {
+            'id': wandb_id,
             'dir': save_dir,
-            'name': args.wandb_exp_name,
+            'name': name,
+            'group': group,
             'project': args.wandb_project,
+            'mode': args.wandb_mode,
+            'resume': 'auto',
             'config': vars(args)}
-        os.makedirs(wandb_kwargs['dir'], exist_ok=True)
+
+        if args.wandb_mode == 'online' or args.wandb_api_key:
+            assert args.wandb_api_key, 'wandb_api_key is required for online mode'
+            wandb.login(key=args.wandb_api_key)
         wandb.init(**wandb_kwargs)
         _GLOBAL_WANDB_WRITER = wandb
 
diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py
index baa6254c..e93636d7 100644
--- a/megatron/training/initialize.py
+++ b/megatron/training/initialize.py
@@ -28,9 +28,12 @@ from megatron.training import get_adlr_autoresume, get_args, get_tensorboard_wri
 from megatron.training.arguments import parse_args, validate_args
 from megatron.training.async_utils import init_persistent_async_worker
 from megatron.training.checkpointing import load_args_from_checkpoint
-from megatron.training.global_vars import set_global_variables
+from megatron.training.global_vars import set_global_variables, set_global_writers
 from megatron.training.yaml_arguments import validate_yaml
 
+from flagscale.train import FSTrainArguments
+from flagscale.train import set_parallel_context, set_get_spiky_loss_detector
+
 logger = logging.getLogger(__name__)
 
 
@@ -80,11 +83,18 @@ def initialize_megatron(
     if args.async_save and args.use_persistent_ckpt_worker:
         init_persistent_async_worker()
 
+    if args.hetero_process_meshes is not None:
+        fs_argument = FSTrainArguments(args)
+        fs_argument.pre_validate_args()
+
     if args.yaml_cfg is not None:
         args = validate_yaml(args, args_defaults)
     else:
         validate_args(args, args_defaults)
 
+    if args.hetero_process_meshes is not None:
+        fs_argument.post_validate_args()
+
     # set global args, build tokenizer, and set adlr-autoresume,
     # tensorboard-writer, and timers.
     set_global_variables(args)
@@ -112,6 +122,9 @@ def initialize_megatron(
         result_rejected_tracker_filename=args.result_rejected_tracker_filename,
     )
 
+    if args.auto_skip_spiky_loss:
+        set_get_spiky_loss_detector(args=args)
+
     # torch.distributed initialization
     def finish_mpu_init():
         args = get_args()
@@ -135,6 +148,9 @@ def initialize_megatron(
 
             MoEAuxLossAutoScaler.set_loss_scale(torch.ones(1, device=torch.cuda.current_device()))
 
+        # Set tensorboard writer and wandb writer.
+        set_global_writers(args)
+
     if skip_mpu_initialization:
         return None
 
@@ -175,7 +191,8 @@ def _compile_dependencies():
     # Compile dataset C++ code.
     # =========================
     # TODO: move this to ninja
-    if torch.distributed.get_rank() == 0:
+    from megatron.core.datasets.utils import is_built_on_zero_rank
+    if is_built_on_zero_rank():
         start_time = time.time()
         print("> compiling dataset index builder ...")
         from megatron.core.datasets.utils import compile_helpers
@@ -330,11 +347,24 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks):
             'timeout': timedelta(minutes=args.distributed_timeout_minutes),
         }
 
+        if args.enable_hetero and args.hetero_use_cpu_communication:
+            # if not all(device_type == args.hetero_device_types[0] for device_type in args.hetero_device_types):
+            #     init_process_group_kwargs['backend'] = 'gloo'
+            init_process_group_kwargs['backend'] = "cpu:gloo"
+        # TODO: @aoyulong the init_process_group will be hanging if the device_id is set
+        # if packaging.version.Version(torch.__version__) >= packaging.version.Version("2.3.0"):
+        #     init_process_group_kwargs['device_id'] = device_id
+
         torch.distributed.init_process_group(**init_process_group_kwargs)
 
     # Set the tensor model-parallel, pipeline model-parallel, and
     # data-parallel communicators.
     if device_count > 0:
+        # Set the parallel context.
+        if args.enable_hetero:
+            set_parallel_context(args)
+            return
+
         if mpu.model_parallel_is_initialized():
             print("model parallel is already initialized")
         else:
diff --git a/megatron/training/mfu_calculate.py b/megatron/training/mfu_calculate.py
new file mode 100644
index 00000000..d74f1536
--- /dev/null
+++ b/megatron/training/mfu_calculate.py
@@ -0,0 +1,189 @@
+# Args ref from https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/configs/config_671B.json
+# modify from https://github.com/feifeibear/DPSKV3MFU/blob/main/dpskv3_flops.py
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+class MFU():
+    def __init__(
+            self,
+            config: TransformerConfig,
+            args
+        ):
+        self.config = config
+        self.hidden_size = self.config.hidden_size
+        self.inter_dim = self.config.ffn_hidden_size
+        # not support  custom moe_inter_dim in megatron 0.8/mixtral, will update it soon in megatron 0.10
+        self.moe_inter_dim = self.config.ffn_hidden_size
+        self.n_layers = self.config.num_layers
+        self.n_heads = self.config.num_attention_heads
+        self.head_dim = self.config.hidden_size // self.config.num_attention_heads
+        self.v_head_dim = self.config.num_query_groups
+
+        self.num_moe_experts = self.config.num_moe_experts
+        self.moe_router_topk = self.config.moe_router_topk
+        self.dtype = self.config.params_dtype
+
+        self.args = args
+        self.swiglu = self.args.swiglu
+        self.bs = self.args.micro_batch_size
+        self.vocab_size = self.args.padded_vocab_size
+        self.seq_len = self.args.seq_length
+
+        #next args for deepseek v3
+        # self.n_dense_layers = 2
+        # self.n_expert_groups = 8
+        # self.n_limited_groups = 4
+        # self.q_lora_rank = 1536
+        # self.kv_lora_rank = 512
+        # self.qk_nope_head_dim = 128
+        # self.qk_rope_head_dim = 64
+        # self.route_scale = 2.5
+        # self.score_func = "sigmoid"
+
+    def cal_embed_fwd_flops(self):
+        #embedding layer
+        return 2 * self.bs * self.seq_len * self.hidden_size
+
+    def cal_head_fwd_flops(self):
+        #lm head layer
+        return 2 * self.bs * self.seq_len * self.hidden_size * self.vocab_size
+
+    def cal_attn_fwd_flops(self):
+        # qkv_proj:  For normal attention without groups, num_query_groups == num_attention_heads, so these two will be the same
+        self.query_projection_size = self.config.kv_channels * self.config.num_attention_heads
+        self.kv_projection_size = self.config.kv_channels * self.config.num_query_groups
+        qkv_proj = 2 * self.bs * self.seq_len * self.hidden_size * (self.query_projection_size + 2 * self.kv_projection_size)
+        # print('qkv_proj', qkv_proj)
+        #GQA and MHA are the same flops, calculate MHA hear
+        flops = 2 * self.bs * self.n_heads * self.seq_len  * self.head_dim * self.seq_len
+        flops += 2 * self.bs * self.n_heads * self.seq_len * self.seq_len  * self.head_dim
+        # print('flops', flops)
+        #flash attn output : self.bs * self.seq_len  * self.hidden_size
+        #weight : self.hidden_size * self.query_projection_size
+        wo_flops = 2 *  self.bs * self.seq_len  * self.hidden_size * self.query_projection_size
+        # print('wo_flops', wo_flops)
+        return qkv_proj + flops / 2 + wo_flops
+
+    def cal_moe_fwd_flops(self):
+        mlp_times = 3 if self.swiglu else 2
+        expert_per_token = 2 * 1 * self.hidden_size * self.moe_inter_dim * mlp_times * self.moe_router_topk
+        flops = self.bs * self.seq_len * expert_per_token
+        return flops
+
+
+
+    def cal_mlp_fwd_flops(self):
+        mlp_times = 3 if self.swiglu else 2 #3 for llama and 2 for gpt
+        flops = 2 * self.bs * self.seq_len * self.hidden_size * self.inter_dim * mlp_times
+        return flops
+
+    def cal_fwd_flops(self):
+        """
+            flops (TFLOPS) per forward
+        """
+        flops_embed = self.cal_embed_fwd_flops()
+        flops_attn = self.cal_attn_fwd_flops()  * self.n_layers
+        if self.num_moe_experts:
+            flops_mlp = self.cal_moe_fwd_flops()  * self.n_layers
+        else:
+            flops_mlp = self.cal_mlp_fwd_flops()  * self.n_layers
+        flops_head = self.cal_head_fwd_flops()
+        flops = flops_embed + flops_attn + flops_mlp  + flops_head
+        # print(f"FWD Flops: {flops} FLOPS,  flops_embed: {flops_embed} FLOPS, flops_attn: {flops_attn} FLOPS, , flops_mlp: {flops_mlp} FLOPS, , flops_head: {flops_head} FLOPS,  ")
+        return flops
+
+    def cal_deepseek_attn_fwd_flops(self):
+        # score = Q x K^T /2 double to causal
+        # scores = torch.einsum("bshd,bthd->bsht", q, self.k_cache[:bsz, :end_pos]) * self.softmax_scale -> [bs, seq_len, seq_len]
+        flops = 2 * self.bs * self.seq_len * self.seq_len * self.n_heads * self.qk_head_dim
+
+        # score x V
+        # x = torch.einsum("bsht,bthd->bshd", scores, self.v_cache[:bsz, :end_pos]) -> [bs, seq_len, self.n_heads, self.v_head_dim]
+        flops += 2 * self.bs * self.seq_len * self.seq_len * self.n_heads * self.v_head_dim
+
+        return flops / 2
+
+    def cal_deepseek_mla_fwd_flops(self):
+        flops = 0
+
+        # 192
+        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+
+        # Q down + up
+        # q = self.wq_b(self.q_norm(self.wq_a(x))) -> [bs, seq_len, (self.n_heads * self.qk_head_dim)]
+        flops += 2 * self.bs * self.seq_len * self.hidden_size * self.q_lora_rank
+        flops += 2 * self.bs * self.seq_len * self.q_lora_rank * self.n_heads * self.qk_head_dim
+
+        # KV down
+        # kv = self.wkv_a(x) -> [bs, seq_len, (self.kv_lora_rank + self.qk_rope_head_dim)]
+        flops += 2 * self.bs * self.seq_len * self.hidden_size * (self.kv_lora_rank + self.qk_rope_head_dim)
+
+        # KV up
+        # kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+        # kv = self.wkv_b(self.kv_norm(kv)) -> [bs, seq_len, (self.n_heads * self.qk_head_dim)]
+        flops += 2 * self.bs * self.seq_len * self.kv_lora_rank * self.n_heads * (self.qk_nope_head_dim + self.v_head_dim)
+
+        # attn
+        flops += self.cal_attn_fwd_flops(self.bs, seq_len)
+
+        # x = self.wo(x.flatten(2))
+        flops += 2 * self.bs * self.seq_len * self.n_heads * self.v_head_dim * self.hidden_size
+
+        return flops
+
+    def cal_deepseek_fwd_flops(self):
+        """
+            flops (TFLOPS) per token
+        """
+
+        shard_expert_num = 1
+        routed_expert_num = 8
+
+        flops_mla = self.cal_mla_fwd_flops() / self.seq_len / self.bs / (BASE**3) * self.n_layers
+        flops_moe = (shard_expert_num + routed_expert_num) * self.cal_moe_fwd_flops() / self.seq_len / self.bs / (BASE**3) * (self.n_layers - self.n_dense_layers)
+        flops_mlp = self.cal_mlp_fwd_flops() / self.seq_len / self.bs / (BASE**3) * self.n_dense_layers
+
+
+        flops_embed = self.cal_embed_fwd_flops() / self.seq_len / self.bs / (BASE**3)
+        flops_head = self.cal_head_fwd_flops() / self.seq_len / self.bs / (BASE**3)
+
+        print(f"flops_mla: {flops_mla} TFLOPS, flops_moe: {flops_moe} TFLOPS")
+
+        flops = flops_mla + flops_moe + flops_mlp + flops_embed + flops_head
+
+        print(f"flops: {flops} TFLOPS")
+        return flops
+
+    def get_mfu(self, tgs, device_tflops):
+        # input:
+            # tgs
+            # tflops for device
+        #output:
+            #mfu
+        fwd_flops = self.cal_fwd_flops()
+        bwd_flops = fwd_flops * 2
+        flops_per_mbs = fwd_flops + bwd_flops
+        Tflops = flops_per_mbs / self.seq_len / self.bs  * tgs / 1e12
+        mfu = Tflops / device_tflops
+        return round(mfu, 4)
+
+if __name__ == '__main__':
+
+    class Args():
+        def __init__(self) -> None:
+            self.swiglu = True
+            self.micro_batch_size = 1
+            self.padded_vocab_size = 32000
+            self.seq_length = 4096
+    args = Args()
+    config = TransformerConfig(
+        num_layers=32,
+        hidden_size=4096,
+        num_attention_heads=32,
+        num_query_groups=8,
+        ffn_hidden_size=14336,
+        num_moe_experts = 8,
+        moe_router_topk = 2
+
+    )
+    mfu = MFU(config, args)
+    print(mfu.get_mfu(1960, 282))
diff --git a/megatron/training/tokenizer/gpt2_tokenization.py b/megatron/training/tokenizer/gpt2_tokenization.py
index 55b95b8e..68e686ec 100644
--- a/megatron/training/tokenizer/gpt2_tokenization.py
+++ b/megatron/training/tokenizer/gpt2_tokenization.py
@@ -322,3 +322,55 @@ class GPT2Tokenizer(object):
                 index += 1
 
         return vocab_file, merge_file, special_tokens_file
+
+
+class AquilaTokenizer(GPT2Tokenizer):
+    def __init__(self, vocab_file, merges_file, errors='replace',
+                 special_tokens=None, max_len=None):
+        super().__init__(vocab_file, merges_file, errors=errors,
+                         special_tokens=special_tokens, max_len=max_len)
+
+        from .tokenization_utils import Trie
+        self.tokens_trie = Trie()
+        if len(self.special_tokens) > 0:
+            for token in self.special_tokens.keys():
+                self.tokens_trie.add(token)
+
+        for k, v in self.special_tokens_decoder.items():
+            self.decoder[k] = v
+            self.encoder[v] = k
+
+    def _tokenize(self, text):
+        """ Tokenize a string. """
+        bpe_tokens = []
+        for token in re.findall(self.pat, text):
+            if sys.version_info[0] == 2:
+                token = ''.join(self.byte_encoder[ord(b)] for b in token)
+            else:
+                token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def tokenize(self, text):
+        tokens = self.tokens_trie.split(text)
+
+        bpe_tokens = []
+        for token in tokens:
+            if not token:
+                continue
+            if token in self.special_tokens:
+                bpe_tokens.append(token)
+            else:
+                bpe_tokens.extend(self._tokenize(token))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = []
+        for token in tokens:
+            if token in self.special_tokens_decoder:
+                text.append(self.special_tokens_decoder[token])
+            else:
+                text.append(self.decoder[token])
+        text = ''.join(text)
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors=self.errors)
+        return text
diff --git a/megatron/training/tokenizer/tokenization_utils.py b/megatron/training/tokenizer/tokenization_utils.py
new file mode 100644
index 00000000..8ec66b77
--- /dev/null
+++ b/megatron/training/tokenizer/tokenization_utils.py
@@ -0,0 +1,167 @@
+import regex as re
+from collections import OrderedDict
+from typing import Any, Dict, List, Optional, Tuple, Union, overload
+
+class Trie:
+    """
+    Trie in Python. Creates a Trie out of a list of words. The trie is used to split on `added_tokens` in one pass
+    Loose reference https://en.wikipedia.org/wiki/Trie
+    """
+
+    def __init__(self):
+        self.data = {}
+
+    def add(self, word: str):
+        if not word:
+            # Prevent empty string
+            return
+        ref = self.data
+        for char in word:
+            ref[char] = char in ref and ref[char] or {}
+            ref = ref[char]
+        ref[""] = 1
+
+    def split(self, text: str) -> List[str]:
+        states = OrderedDict()
+
+        # This will contain every indices where we need
+        # to cut.
+        # We force to cut at offset 0 and len(text) (added later)
+        offsets = [0]
+
+        # This is used by the lookahead which needs to skip over
+        # some text where the full match exceeded the place in the initial
+        # for loop
+        skip = 0
+        # Main loop, Giving this algorithm O(n) complexity
+        for current, current_char in enumerate(text):
+            if skip and current < skip:
+                # Prevents the lookahead for matching twice
+                # like extra_id_100 and id_100
+                continue
+
+            # This will track every state
+            # that stop matching, we need to stop tracking them.
+            # If we look at "lowball", we're going to match "l" (add it to states), "o", "w", then
+            # fail on "b", we need to remove 0 from the valid states.
+            to_remove = set()
+            # Whenever we found a match, we need to drop everything
+            # this is a greedy algorithm, it will match on the first found token
+            reset = False
+
+            # In this case, we already have partial matches (But unfinished)
+            for start, trie_pointer in states.items():
+                if "" in trie_pointer:
+                    # This is a final match, we need to reset and
+                    # store the results in `offsets`.
+
+                    # Lookahead to match longest first
+                    # Important in case of extra_id_1 vs extra_id_100
+                    # Here we are also actively looking for other earlier partial
+                    # matches
+                    # "[CLS]", "L", we need to match CLS even if L is special
+                    for lookstart, looktrie_pointer in states.items():
+                        if lookstart > start:
+                            # This partial match is later, we can stop looking
+                            break
+                        elif lookstart < start:
+                            # This partial match is earlier, the trie pointer
+                            # was already updated, so index is + 1
+                            lookahead_index = current + 1
+                            end = current + 1
+                        else:
+                            # Here lookstart == start and
+                            #      looktrie_pointer == trie_pointer
+                            # It wasn't updated yet so indices are current ones
+                            lookahead_index = current
+                            end = current
+                        next_char = text[lookahead_index] if lookahead_index < len(text) else None
+                        if "" in looktrie_pointer:
+                            start = lookstart
+                            end = lookahead_index
+                            skip = lookahead_index
+
+                        while next_char in looktrie_pointer:
+                            looktrie_pointer = looktrie_pointer[next_char]
+                            lookahead_index += 1
+                            if "" in looktrie_pointer:
+                                start = lookstart
+                                end = lookahead_index
+                                skip = lookahead_index
+
+                            if lookahead_index == len(text):
+                                # End of string
+                                break
+                            next_char = text[lookahead_index]
+                        # End lookahead
+
+                    # Storing and resetting
+                    offsets.append(start)
+                    offsets.append(end)
+                    reset = True
+                    break
+                elif current_char in trie_pointer:
+                    # The current character being looked at has a match within the trie
+                    # update the pointer (it will be stored back into states later).
+                    trie_pointer = trie_pointer[current_char]
+
+                    # Storing back the new pointer into the states.
+                    # Partial matches got longer by one.
+                    states[start] = trie_pointer
+                else:
+                    # The new character has not match in the trie, we need
+                    # to stop keeping track of this partial match.
+                    # We can't do it directly within the loop because of how
+                    # python iteration works
+                    to_remove.add(start)
+
+            # Either clearing the full start (we found a real match)
+            # Or clearing only the partial matches that didn't work.
+            if reset:
+                states = {}
+            else:
+                for start in to_remove:
+                    del states[start]
+
+            # If this character is a starting character within the trie
+            # start keeping track of this partial match.
+            if current >= skip and current_char in self.data:
+                states[current] = self.data[current_char]
+
+        # We have a cut at the end with states.
+        for start, trie_pointer in states.items():
+            if "" in trie_pointer:
+                # This is a final match, we need to reset and
+                # store the results in `offsets`.
+                end = len(text)
+                offsets.append(start)
+                offsets.append(end)
+                # Longest cut is always the one with lower start so the first
+                # item so we need to break.
+                break
+
+        return self.cut_text(text, offsets)
+
+    def cut_text(self, text, offsets):
+        # We have all the offsets now, we just need to do the actual splitting.
+        # We need to eventually add the first part of the string and the eventual
+        # last part.
+        offsets.append(len(text))
+        tokens = []
+        start = 0
+        for end in offsets:
+            if start > end:
+                logger.error(
+                    "There was a bug in Trie algorithm in tokenization. Attempting to recover. Please report it"
+                    " anyway."
+                )
+                continue
+            elif start == end:
+                # This might happen if there's a match at index 0
+                # we're also preventing zero-width cuts in case of two
+                # consecutive matches
+                continue
+            tokens.append(text[start:end])
+            start = end
+
+        return tokens
\ No newline at end of file
diff --git a/megatron/training/tokenizer/tokenizer.py b/megatron/training/tokenizer/tokenizer.py
index e728d91f..f66f199f 100644
--- a/megatron/training/tokenizer/tokenizer.py
+++ b/megatron/training/tokenizer/tokenizer.py
@@ -14,6 +14,7 @@ from megatron.core.datasets.megatron_tokenizer import MegatronTokenizer
 
 from .bert_tokenization import FullTokenizer as FullBertTokenizer
 from .gpt2_tokenization import GPT2Tokenizer
+from .gpt2_tokenization import AquilaTokenizer
 from megatron.training.tokenizer.multimodal_tokenizer import MultimodalTokenizer
 
 
@@ -96,6 +97,31 @@ def build_tokenizer(args, **kwargs):
     elif args.tokenizer_type == 'NullMultimodalTokenizer':
         assert args.vocab_size is not None
         tokenizer = _NullMultimodalTokenizer(args.vocab_size)
+    elif args.tokenizer_type == 'AquilaTokenizerFS':
+        assert args.vocab_file is not None
+        assert args.merge_file is not None
+        assert args.special_tokens_file is not None
+        tokenizer = _AquilaTokenizerFS(args.vocab_file, args.merge_file,
+                                     args.special_tokens_file)
+    elif args.tokenizer_type == "HFTokenizerFS":
+        assert args.tokenizer_path is not None
+        tokenizer = _HFTokenizerFS(args.tokenizer_path)
+    elif args.tokenizer_type == "Llama3TokenizerFS":
+        assert args.tokenizer_path is not None
+        tokenizer = _Llama3TokenizerFS(args.tokenizer_path)
+    elif args.tokenizer_type == "QwenTokenizerFS":
+        assert args.tokenizer_path is not None
+        tokenizer = _QwenTokenizerFS(args.tokenizer_path)
+    elif args.tokenizer_type == "HFTokenizersTokenizerFS":
+        assert args.tokenizer_path is not None
+        tokenizer = _HFTokenizersTokenizerFS(args.tokenizer_path)
+    elif args.tokenizer_type == "Qwen2TokenizerFS":
+        assert args.tokenizer_path is not None
+        tokenizer = _Qwen2TokenizerFS(args.tokenizer_path, args)
+    elif args.tokenizer_type == 'Qwen2VLTokenizer':
+        assert args.tokenizer_path is not None
+        tokenizer = _Qwen2VLTokenizer(args.tokenizer_path, args.extra_vocab_size)
+        args.padded_vocab_size = tokenizer.vocab_size # no padding
     else:
         raise NotImplementedError('{} tokenizer is not ' 'implemented.'.format(args.tokenizer_type))
 
@@ -586,6 +612,16 @@ class _Llama2Tokenizer(_SentencePieceTokenizer):
             t = t + [self.eos_id]
         return t
 
+    def instruct_tokenize(self, s: str, bos=True, eos=False):
+        '''Default args for text completion, not chat/dialog.'''
+        assert type(s) is str
+        t = self.tokenizer.encode(s)
+        if bos:
+            t = [self.bos_id] + t
+        if eos:
+            t = t + [self.eos_id]
+        return t
+
     def detokenize(self, ids):
         return self.tokenizer.decode_ids(ids)
 
@@ -899,3 +935,273 @@ class _NullMultimodalTokenizer(MegatronTokenizer):
     @property
     def additional_special_tokens_ids(self):
         return None
+
+
+class _AquilaTokenizerFS(MegatronTokenizer):
+    """Aquila tokenizer."""
+
+    def __init__(self, vocab_file, merge_file, special_tokens_file):
+        super().__init__(vocab_file, merge_file, special_tokens_file)
+
+        special_tokens = []
+        if special_tokens_file:
+            special_tokens = open(special_tokens_file, encoding='utf-8').read().split('\n')[:-1]
+
+        self.tokenizer = AquilaTokenizer(vocab_file, merge_file, errors='replace',
+                                         special_tokens=special_tokens, max_len=None)
+        self.eod_id = self.tokenizer.encoder['</s>']
+        self.cls_id = self.tokenizer.encoder['[CLS]']
+        self.pad_id = self.tokenizer.encoder['<|endoftext|>']
+
+    @property
+    def vocab_size(self):
+        return len(self.tokenizer.encoder)
+
+    @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.eod_id
+
+    @property
+    def cls(self):
+        return self.cls_id
+
+    @property
+    def pad(self):
+        return self.pad_id
+
+
+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 _Llama3TokenizerFS(_HFTokenizerFS):
+
+    def __init__(self, tokenizer_path):
+        super().__init__(tokenizer_path)
+
+    @property
+    def vocab_size(self):
+        return self.tokenizer.vocab_size + len(self.tokenizer.get_added_vocab())
+
+
+class _QwenTokenizerFS(_HFTokenizerFS):
+    """Adapted Qwen tokenizer."""
+
+    def __init__(self, tokenizer_path):
+        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]
+
+
+class _HFTokenizersTokenizerFS(MegatronTokenizer):
+    """Tokenizer from HuggingFace Tokenizers."""
+
+    def __init__(self, json_file):
+        super().__init__(json_file)
+
+        from tokenizers import Tokenizer
+        self.tokenizer = Tokenizer.from_file(json_file)
+
+        print(f"Vocab size: {self.tokenizer.get_vocab_size()}")
+
+        self.eod_id = self.tokenizer.token_to_id("<|endoftext|>")
+        self.pad_id = self.tokenizer.token_to_id("<|padding|>")
+
+        self._inv_vocab = None
+
+    @property
+    def vocab_size(self):
+        return self.tokenizer.get_vocab_size()
+
+    @property
+    def vocab(self):
+        return self.tokenizer.get_vocab()
+
+    @property
+    def inv_vocab(self):
+        # return self.tokenizer.decoder
+        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 pad(self):
+        return self.pad_id
+
+
+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
+
+
+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)
\ No newline at end of file
diff --git a/megatron/training/utils.py b/megatron/training/utils.py
index 61a69e14..50848406 100644
--- a/megatron/training/utils.py
+++ b/megatron/training/utils.py
@@ -51,6 +51,7 @@ try:
 except ImportError:
     ALL_MODULE_WRAPPER_CLASSNAMES = (DDP, custom_FSDP, Float16Module)
 
+from flagscale.train.hetero.p2p_communication import get_device_type_for_comm
 
 def unwrap_model(model, module_instances=ALL_MODULE_WRAPPER_CLASSNAMES):
     return_list = True
@@ -177,35 +178,71 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False):
     else:
         moe_norm_2 = torch.zeros_like(norm_2)
 
-    # Reduce norm across model parallel groups (dense and expert).
-    # Dense params should sum across all model-parallel GPUs (tensor + pipeline).
-    dense_reduce_group = mpu.get_model_parallel_group()
-    ranks_in_dense_reduce_group = torch.distributed.get_process_group_ranks(dense_reduce_group)
-    # Expert params should sum across all model-parallel GPUs (expert + tensor + pipeline).
-    expert_reduce_group = mpu.get_expert_tensor_model_pipeline_parallel_group()
-    ranks_in_expert_reduce_group = torch.distributed.get_process_group_ranks(expert_reduce_group)
-
-    # If dense and expert reduce groups are the same, sum then reduce.
-    if ranks_in_dense_reduce_group == ranks_in_expert_reduce_group:
-        norm_2 += moe_norm_2
-        torch.distributed.all_reduce(
-            norm_2,
-            op=torch.distributed.ReduceOp.SUM,
-            group=dense_reduce_group
-        )
-    # If dense and expert reduce groups are different, reduce then sum.
-    else:
-        torch.distributed.all_reduce(
-            norm_2,
-            op=torch.distributed.ReduceOp.SUM,
-            group=dense_reduce_group
-        )
-        torch.distributed.all_reduce(
-            moe_norm_2,
-            op=torch.distributed.ReduceOp.SUM,
-            group=expert_reduce_group
-        )
-        norm_2 += moe_norm_2
+    ########## FlagScale Begin ##########
+    # Sum across all model-parallel GPUs(tensor + pipeline).
+    mp_groups = mpu.get_model_parallel_group()
+    comm_device = get_device_type_for_comm(mp_groups)
+    if comm_device == "cpu":
+        norm_2 = norm_2.cpu()
+    if isinstance(mp_groups, list):  # hetero
+        original_norm_2 = norm_2.clone().detach()
+        for mp_group in mp_groups:
+            norm_2.copy_(original_norm_2)
+            torch.distributed.all_reduce(
+                norm_2, op=torch.distributed.ReduceOp.SUM, group=mp_group
+            )
+        if len(moe_params_data) > 0:
+            emp_groups = mpu.get_expert_tensor_model_pipeline_parallel_group()
+            comm_device = get_device_type_for_comm(emp_groups)
+            if comm_device == "cpu":
+                moe_norm_2 = moe_norm_2.cpu()
+
+            assert isinstance(
+                emp_groups, list
+            ), "emp_groups should be a list if mp_groups is a list"
+            original_norm_2 = moe_norm_2.clone().detach()
+            for emp_group in emp_groups:
+                moe_norm_2.copy_(original_norm_2)
+                torch.distributed.all_reduce(
+                    moe_norm_2, op=torch.distributed.ReduceOp.SUM, group=emp_group
+                )
+            norm_2 += moe_norm_2
+    ########## FlagScale End ##########
+    else:  # original code
+
+        # Reduce norm across model parallel groups (dense and expert).
+        # Dense params should sum across all model-parallel GPUs (tensor + pipeline).
+        dense_reduce_group = mpu.get_model_parallel_group()
+        ranks_in_dense_reduce_group = torch.distributed.get_process_group_ranks(dense_reduce_group)
+        # Expert params should sum across all model-parallel GPUs (expert + tensor + pipeline).
+        expert_reduce_group = mpu.get_expert_tensor_model_pipeline_parallel_group()
+        ranks_in_expert_reduce_group = torch.distributed.get_process_group_ranks(expert_reduce_group)
+
+        # If dense and expert reduce groups are the same, sum then reduce.
+        if ranks_in_dense_reduce_group == ranks_in_expert_reduce_group:
+            norm_2 += moe_norm_2
+            torch.distributed.all_reduce(
+                norm_2,
+                op=torch.distributed.ReduceOp.SUM,
+                group=dense_reduce_group
+            )
+        # If dense and expert reduce groups are different, reduce then sum.
+        else:
+            torch.distributed.all_reduce(
+                norm_2,
+                op=torch.distributed.ReduceOp.SUM,
+                group=dense_reduce_group
+            )
+            torch.distributed.all_reduce(
+                moe_norm_2,
+                op=torch.distributed.ReduceOp.SUM,
+                group=expert_reduce_group
+            )
+            norm_2 += moe_norm_2
+
+    if comm_device == "cpu":
+        norm_2 = norm_2.cuda()
+        moe_norm_2 = moe_norm_2.cuda()
 
     return norm_2.item() ** 0.5
 
@@ -232,10 +269,18 @@ def reduce_max_stat_across_model_parallel_group(stat: float) -> float:
     """
     if stat is None:
         stat = -1.0
-    stat = torch.tensor([stat], dtype=torch.float32, device=torch.cuda.current_device())
-    torch.distributed.all_reduce(
-        stat, op=torch.distributed.ReduceOp.MAX, group=mpu.get_model_parallel_group()
-    )
+    model_parallel_groups = mpu.get_model_parallel_group()
+    if not isinstance(model_parallel_groups, list):
+        stat = torch.tensor([stat], dtype=torch.float32, device=get_device_type_for_comm(model_parallel_groups))
+        torch.distributed.all_reduce(
+            stat, op=torch.distributed.ReduceOp.MAX, group=mpu.get_model_parallel_group()
+        )
+    else:
+        stat = torch.tensor([stat], dtype=torch.float32, device=get_device_type_for_comm(model_parallel_groups[0]))
+        for model_parallel_group in model_parallel_groups:
+            torch.distributed.all_reduce(
+                stat, op=torch.distributed.ReduceOp.MAX, group=model_parallel_group
+            )
     if stat.item() == -1.0:
         return None
     else:
@@ -250,10 +295,18 @@ def logical_and_across_model_parallel_group(input: bool) -> bool:
         input = 1
     else:
         input = 0
-    input = torch.tensor([input], dtype=torch.int, device=torch.cuda.current_device())
-    torch.distributed.all_reduce(
-        input, op=torch.distributed.ReduceOp.MIN, group=mpu.get_model_parallel_group()
-    )
+    model_parallel_groups = mpu.get_model_parallel_group()
+    if not isinstance(model_parallel_groups, list):
+        input = torch.tensor([input], dtype=torch.int, device=get_device_type_for_comm(model_parallel_groups))
+        torch.distributed.all_reduce(
+            input, op=torch.distributed.ReduceOp.MIN, group=mpu.get_model_parallel_group()
+        )
+    else:
+        input = torch.tensor([input], dtype=torch.int, device=get_device_type_for_comm(model_parallel_groups[0]))
+        for model_parallel_group in model_parallel_groups:
+            torch.distributed.all_reduce(
+                input, op=torch.distributed.ReduceOp.MIN, group=model_parallel_group
+            )
     return bool(input.item())
 
 
@@ -269,7 +322,13 @@ def report_memory(name):
         torch.cuda.memory_reserved() / mega_bytes)
     string += ' | max reserved: {}'.format(
         torch.cuda.max_memory_reserved() / mega_bytes)
-    if mpu.get_data_parallel_rank() == 0:
+    args = get_args()
+    if not args.auto_tune:
+        # Each rank prints the memory report.
+        if mpu.get_data_parallel_rank() == 0:
+            print("[Rank {}] {}".format(torch.distributed.get_rank(), string),
+                  flush=True)
+    else:
         print("[Rank {}] {}".format(torch.distributed.get_rank(), string),
               flush=True)
 
@@ -384,8 +443,11 @@ def is_rank0():
     return torch.distributed.is_initialized() and torch.distributed.get_rank() == 0
 
 def is_last_rank():
-    return torch.distributed.get_rank() == (
-        torch.distributed.get_world_size() - 1)
+    if mpu.get_pipeline_model_parallel_world_size() > 1:
+        return torch.distributed.get_rank() == mpu.get_last_rank_when_using_pipeline()
+    else:
+        return torch.distributed.get_rank() == (
+            torch.distributed.get_world_size() - 1)
 
 def print_rank_last(message):
     """If distributed is initialized, print only on last rank."""
diff --git a/tasks/aquila/datasets.py b/tasks/aquila/datasets.py
new file mode 100644
index 00000000..d0e82d44
--- /dev/null
+++ b/tasks/aquila/datasets.py
@@ -0,0 +1,75 @@
+"""Aquila datasets."""
+
+import json
+import math
+
+import numpy as np
+import torch
+
+from megatron import get_args
+from megatron import print_rank_0
+from megatron import get_tokenizer
+
+
+def build_dataset(task):
+    """Helper function to select and build dataset."""
+
+    if task == 'AQUILA':
+        return _build_aquila_dataset()
+
+    raise NotImplementedError('dataset for {} task is not '
+                              'implemented.'.format(task))
+
+
+class _AquilaDataset(torch.utils.data.Dataset):
+
+    def __init__(self, path, tokenizer, seq_len):
+        print_rank_0('> building aquila dataset from {} ...'.format(path))
+        self.seq_len = seq_len
+        self.tokenizer = tokenizer
+        self.BOS_TOKEN = self.tokenizer.cls
+        self.EOS_TOKEN = self.tokenizer.eod
+        # 2048 for 7B
+        self.text_maxlen = seq_len
+
+        import jsonlines
+        self.texts = []
+        with jsonlines.open(path) as reader:
+            for line in reader:
+                if 'text' not in line:
+                    continue
+                text = line['text'][:self.text_maxlen]
+                self.texts.append(text)
+
+    def __len__(self):
+        return len(self.texts)
+
+    def __getitem__(self, idx):
+        text = self.texts[idx]
+        tokens = [self.BOS_TOKEN]
+        tokens += self.tokenizer.tokenize(text)
+        tokens.append(self.EOS_TOKEN)
+        tokens = tokens[:self.seq_len+1]
+        num_tokens = len(tokens)
+        pad_mask = [1] * num_tokens
+        if num_tokens < self.seq_len + 1:
+            num_pad = (self.seq_len + 1 - num_tokens)
+            pad_mask += [0] * (num_pad)
+            tokens += [0] * (num_pad)
+        pad_mask = np.array(pad_mask[1:])
+        tokens = np.array(tokens)
+
+        return {'text': tokens, 'pad_mask': pad_mask}
+
+
+def _build_aquila_dataset():
+    """Build aquila dataset."""
+    args = get_args()
+    tokenizer = get_tokenizer()
+
+    assert len(args.valid_data) == 1
+    val_dataset = _AquilaDataset(args.valid_data[0], tokenizer,
+                                 args.seq_length)
+    print_rank_0(' > found {} samples.'.format(len(val_dataset)))
+
+    return val_dataset
diff --git a/tasks/aquila/evaluate.py b/tasks/aquila/evaluate.py
new file mode 100644
index 00000000..9465b7c4
--- /dev/null
+++ b/tasks/aquila/evaluate.py
@@ -0,0 +1,210 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+"""GPT zero-shot evaluation."""
+
+import math
+
+import torch
+
+from megatron import get_args
+from megatron import print_rank_0, is_last_rank
+from megatron import get_tokenizer
+from megatron.core import parallel_state, tensor_parallel
+from megatron.checkpointing import load_checkpoint
+from megatron.model import GPTModel
+from megatron.training import get_model
+from megatron.utils import get_ltor_masks_and_position_ids, unwrap_model
+from megatron.core.pipeline_parallel.p2p_communication import recv_forward, send_forward
+from megatron.arguments import core_transformer_config_from_args
+from tasks.finetune_utils import build_data_loader
+
+from .datasets import build_dataset
+
+
+def get_model_provider(eval_metric):
+    """Based on evaluation metric set the parallel-output flag and
+    return the model provider."""
+
+    def model_provider(pre_process=True, post_process=True):
+        """Build the model."""
+
+        config = core_transformer_config_from_args(get_args())
+
+        if eval_metric == 'loss':
+            parallel_output = True
+        elif eval_metric == 'accuracy':
+            parallel_output = False
+        else:
+            raise NotImplementedError('output type for {} evaluation metric '
+                                      'is not supported.'.format(eval_metric))
+
+        print_rank_0('building GPT model ...')
+        model = GPTModel(config, num_tokentypes=0, parallel_output=parallel_output,
+                         pre_process=pre_process, post_process=post_process)
+
+        return model
+
+    return model_provider
+
+
+def process_batch(batch):
+    """Process batch and produce inputs for the model."""
+    args = get_args()
+    tokenizer = get_tokenizer()
+
+    loss_mask = batch['pad_mask'].long().cuda().contiguous().byte()
+    tokens_ = batch['text'].long().cuda().contiguous()
+    labels = tokens_[:, 1:].contiguous()
+    tokens = tokens_[:, :-1].contiguous()
+
+    # Get the masks and postition ids.
+    attention_mask, _, position_ids = get_ltor_masks_and_position_ids(
+        tokens,
+        tokenizer.eod,
+        args.reset_position_ids,
+        args.reset_attention_mask,
+        args.eod_mask_loss)
+
+    return tokens, labels, attention_mask, position_ids, loss_mask
+
+
+def forward_step(batch, model, eval_metric, config):
+    """Forward step."""
+
+    # Get the batch.
+    tokens, labels, attention_mask, position_ids, loss_mask = process_batch(
+        batch)
+
+    # Tell the model what our actual batch size will be
+    args = get_args()
+    args.micro_batch_size = len(labels)
+
+    tensor_shape = (args.seq_length, args.micro_batch_size, args.hidden_size)
+    input_tensor = recv_forward(tensor_shape, config)
+
+    # Forward pass through the model.
+    unwrapped_model = unwrap_model(model)
+    unwrapped_model.set_input_tensor(input_tensor)
+    output = model(tokens, position_ids, attention_mask)
+
+    send_forward(output, config)
+
+    if parallel_state.is_pipeline_last_stage():
+        # For loss, return the unreduced loss.
+        if eval_metric == 'loss':
+            losses = tensor_parallel.vocab_parallel_cross_entropy(
+                output.contiguous().float(), labels.contiguous())
+            loss = torch.sum(
+                losses.view(-1) * loss_mask.contiguous().view(-1).float())
+            loss_mask = torch.sum(
+                loss_mask.contiguous().view(-1).float())
+            return loss / loss_mask
+
+        # For accuracy, return the number of correctly predicted samples.
+        if eval_metric == 'accuracy':
+            outputs = torch.argmax(output, -1)
+            correct = (outputs == labels).float()
+            correct[(1 - loss_mask).bool()] = 1
+            correct = correct.prod(-1)
+            return correct.sum()
+
+        raise NotImplementedError('forward method for evaluation metric {} '
+                                  'is not implemented.'.format(eval_metric))
+    return None
+
+
+def evaluate(data_loader, model, eval_metric):
+    """Evaluation."""
+    args = get_args()
+    config = core_transformer_config_from_args(args)
+
+    # Turn on evaluation mode which disables dropout.
+    model.eval()
+
+    total_output = 0.0
+    with torch.no_grad():
+        # For all the batches in the dataset.
+        for iteration, batch in enumerate(data_loader):
+            if iteration % args.log_interval == 0:
+                print_rank_0('> working on iteration: {}'.format(iteration))
+            # Forward evaluation.
+            output = forward_step(batch, model, eval_metric, config)
+
+            # Reduce across processes.
+            if parallel_state.is_pipeline_last_stage():
+                torch.distributed.all_reduce(output,
+                                             group=parallel_state.get_data_parallel_group())
+
+                total_output += output
+
+    return total_output
+
+
+def evaluate_and_print_results(task, data_loader, model, eval_metric):
+    """Evaluate and print results on screen."""
+
+    # Evaluate and get results.
+    output = evaluate(data_loader, model, eval_metric)
+
+    string = ' validation results on {} | '.format(task)
+    if is_last_rank():
+        if eval_metric == 'loss':
+            num_tokenized_tokens = data_loader.dataset.num_tokenized_tokens
+            num_original_tokens = data_loader.dataset.num_original_tokens
+            val_loss = output / (num_tokenized_tokens - 1)
+            ppl = math.exp(min(20, val_loss))
+            token_ratio = (num_tokenized_tokens - 1) / (num_original_tokens - 1)
+            adjusted_ppl = math.exp(min(20, val_loss * token_ratio))
+            string += 'avg loss: {:.4E} | '.format(val_loss)
+            string += 'ppl: {:.4E} | '.format(ppl)
+            string += 'adjusted ppl: {:.4E} | '.format(adjusted_ppl)
+            string += 'token ratio: {} |'.format(token_ratio)
+
+        elif eval_metric == 'accuracy':
+            num_examples = len(data_loader.dataset)
+            acc = output / num_examples
+            string += 'number correct: {:.4E} | '.format(output)
+            string += 'total examples: {:.4E} | '.format(num_examples)
+            string += 'avg accuracy: {:.4E}'.format(acc)
+
+        else:
+            raise NotImplementedError('evaluation method for {} metric is not '
+                                      'implemented yet.'.format(eval_metric))
+
+        length = len(string) + 1
+        print('-' * length)
+        print(string)
+        print('-' * length)
+
+
+def main():
+    """Main program."""
+    args = get_args()
+
+    if args.num_layers_per_virtual_pipeline_stage is not None:
+        print("Interleaved pipeline schedule is not yet supported for evaluation.")
+        exit()
+
+    if args.eval_metric is not None:
+        eval_metric = args.eval_metric
+    else:
+        raise NotImplementedError('{} task is not implemented.'.format(
+            args.task))
+
+    # Set up model and load checkpoint.
+    model = get_model(get_model_provider(eval_metric), wrap_with_ddp=False)
+    if args.load is not None:
+        _ = load_checkpoint(model, None, None)
+
+    assert len(model) == 1, "Above condition should have caught this"
+    model = model[0]
+
+    # Data stuff.
+    dataset = build_dataset(args.task)
+    dataloader = build_data_loader(dataset, args.micro_batch_size,
+                                   args.num_workers, drop_last=False)
+
+    # Run evaluation.
+    evaluate_and_print_results(args.task, dataloader, model, eval_metric)
+
+    print_rank_0('done :-)')
diff --git a/tasks/main.py b/tasks/main.py
index da8c4b9b..9bba4424 100644
--- a/tasks/main.py
+++ b/tasks/main.py
@@ -32,6 +32,9 @@ def get_tasks_args(parser):
                        help='Sliding window for overlapping evaluation.')
     group.add_argument('--strict-lambada', action='store_true',
                        help='Use more difficult formulation of lambada.')
+    group.add_argument('--eval-metric', type=str, default='loss',
+                       choices=['loss', 'accuracy'],
+                       help='Metric for evaluation tasks')
     # Retriever args
     group.add_argument('--qa-data-dev', type=str, default=None,
                        help='Path to the QA dataset dev file.')
@@ -89,6 +92,8 @@ if __name__ == '__main__':
         from glue.finetune import main
     elif args.task in ['LAMBADA', 'WIKITEXT103']:
         from zeroshot_gpt.evaluate import main
+    elif args.task in ['AQUILA']:
+        from aquila.evaluate import main
     elif args.task in ['ICT-ZEROSHOT-NQ', 'RETRIEVER-EVAL']:
         from orqa.evaluate_orqa import main
     elif args.task in ['RET-FINETUNE-NQ']:
diff --git a/tools/checkpoint/saver_legacy.py b/tools/checkpoint/saver_legacy.py
index 50af6a57..ab8cfebb 100644
--- a/tools/checkpoint/saver_legacy.py
+++ b/tools/checkpoint/saver_legacy.py
@@ -132,6 +132,8 @@ def save_checkpoint(queue, args):
         sys.argv.append('--untie-embeddings-and-output-weights')
     if not md.linear_bias:
         sys.argv.append('--disable-bias-linear')
+    if not md.linear_bias_qkv:
+        sys.argv.append('--disable-bias-linear-qkv')
 
     if md.model_type == 'BERT' and not md.bert_binary_head:
         sys.argv.append('--bert-no-binary-head')
@@ -359,7 +361,35 @@ def save_checkpoint(queue, args):
                 if not hasattr(models[0].language_model, 'output_layer'):
                     print("ERROR: got an output layer, but model does not have one")
                     exit(1)
-                output_layer_weight = torch.chunk(msg.pop("weight"), args.target_tensor_parallel_size, dim=0)
+                # Deal with padding
+                orig_output_layer_weight = msg.pop("weight")
+                if md.true_vocab_size is not None:
+                    # figure out what our padded vocab size is
+                    orig_output_layer_size = orig_output_layer_weight.shape[0]
+                    margs.padded_vocab_size = _vocab_size_with_padding(md.true_vocab_size, margs)
+
+                    # Cut out extra padding we don't need
+                    if orig_output_layer_size > margs.padded_vocab_size:
+                        full_output_layer_weight = orig_output_layer_weight[0:margs.padded_vocab_size,:]
+
+                    # Expanding embedding to larger size by replicating final entry
+                    elif orig_output_layer_size < margs.padded_vocab_size:
+                        padding_size = margs.padded_vocab_size - orig_output_layer_size
+
+                        full_output_layer_weight = torch.cat((
+                            orig_output_layer_weight,
+                            orig_output_layer_weight[-1].unsqueeze(0).expand(padding_size, -1)))
+
+                    # Same size!
+                    else:
+                        full_output_layer_weight = orig_output_layer_weight
+                else:
+                    print("Original vocab size not specified, leaving embedding table as-is. "
+                          "If you've changed the tensor parallel size this could cause problems.")
+                    margs.padded_vocab_size = orig_output_layer_weight.shape[0]
+                    full_output_layer_weight = orig_output_layer_weight
+
+                output_layer_weight = torch.chunk(full_output_layer_weight, args.target_tensor_parallel_size, dim=0)
                 for tp_rank in range(args.target_tensor_parallel_size):
                     models[tp_rank].language_model.output_layer.weight.data.copy_(output_layer_weight[tp_rank])
                 del output_layer_weight
diff --git a/tools/preprocess_data.py b/tools/preprocess_data.py
index 13e5b64a..c90feb5c 100644
--- a/tools/preprocess_data.py
+++ b/tools/preprocess_data.py
@@ -12,6 +12,7 @@ import time
 import gzip
 import glob
 import torch
+import shutil
 import numpy as np
 import multiprocessing
 try:
@@ -258,7 +259,12 @@ def main():
 
     if args.split_sentences:
         if nltk_available:
-            nltk.download("punkt", quiet=True, download_dir=os.environ.get("NLTK_DATA"))
+            try:
+                punkt_path = os.environ.get("NLTK_DATA") + "/tokenizers/punkt"
+                if not os.path.exists(punkt_path):
+                    shutil.copytree('/root/nltk_data/tokenizers/punkt', punkt_path)
+            except:
+                nltk.download("punkt", quiet=True, download_dir=os.environ.get("NLTK_DATA"))
         else:
             raise Exception(
                 "nltk library required for sentence splitting is not available.")
