diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py
index 7e94a0be111..8c68399dfe6 100644
--- a/examples/post_training/modelopt/quantize.py
+++ b/examples/post_training/modelopt/quantize.py
@@ -344,7 +344,7 @@ def get_calib_dataloader(
 
     Supports either a local path (.jsonl) or a HuggingFace dataset name.
     """
-    if os.path.isfile(dataset_path_or_name):
+    if os.path.isfile(dataset_path_or_name) and dataset_path_or_name.endswith(".jsonl"):
         # Local file
         print_rank_0(f"Loading calibration dataset from local file: {dataset_path_or_name}")
         all_texts = []
diff --git a/megatron/post_training/checkpointing.py b/megatron/post_training/checkpointing.py
index 1e631b54932..1cb730bc450 100644
--- a/megatron/post_training/checkpointing.py
+++ b/megatron/post_training/checkpointing.py
@@ -8,12 +8,16 @@
 import modelopt
 import modelopt.torch.opt as mto
 import torch.nn as nn
-from modelopt.torch.opt.plugins import restore_sharded_modelopt_state as restore_sharded_modelopt_state_legacy
-from modelopt.torch.opt.plugins.mcore_dist_checkpointing import _load_extra_state_from_sharded_checkpoint
+from modelopt.torch.opt.plugins import (
+    restore_sharded_modelopt_state as restore_sharded_modelopt_state_legacy,
+)
+from modelopt.torch.opt.plugins.mcore_dist_checkpointing import (
+    _load_extra_state_from_sharded_checkpoint,
+)
 
 from megatron.core import dist_checkpointing
-from megatron.core.utils import get_torch_version, is_torch_min_version, unwrap_model
 from megatron.core.dist_checkpointing.serialization import _legacy_common_state_exists
+from megatron.core.utils import unwrap_model
 from megatron.training import get_args
 from megatron.training.checkpointing import _load_base_checkpoint, load_checkpoint
 from megatron.training.utils import print_rank_0
@@ -226,3 +230,24 @@ def restore_sharded_modelopt_state(model: list[nn.Module], checkpoint_name: str
 
     model[0] = mto.restore_from_modelopt_state(model[0], common_modelopt_state)
     _load_extra_state_from_sharded_checkpoint(model[0], checkpoint_name, prefix="")
+
+
+def load_kd_teacher_checkpoint(model) -> None:
+    """Load the teacher checkpoint for ModelOpt distillation if the model has one."""
+    args = get_args()
+    if not getattr(args, "export_kd_teacher_load", None):
+        return
+
+    teacher = unwrap_model(model[0]).teacher_model
+    print_rank_0(f"Loading teacher as {type(teacher).__name__} from {args.export_kd_teacher_load} ...")
+    # [WAR]: To avoid error out on loading teacher's checkpoint, we temporarily
+    # set args.finetune to True while loading the teacher checkpoint.
+    original_args_finetune, original_ckpt_format = args.finetune, args.ckpt_format
+    args.finetune = True
+    if args.export_kd_teacher_ckpt_format is not None:
+        args.ckpt_format = args.export_kd_teacher_ckpt_format
+    try:
+        load_checkpoint([teacher], None, None, load_arg='export_kd_teacher_load')
+    finally:
+        args.finetune, args.ckpt_format = original_args_finetune, original_ckpt_format
+    print_rank_0("... teacher loaded successfully.")
diff --git a/megatron/post_training/model_builder.py b/megatron/post_training/model_builder.py
index 0b411788115..cd497907406 100644
--- a/megatron/post_training/model_builder.py
+++ b/megatron/post_training/model_builder.py
@@ -5,11 +5,11 @@
 import logging
 import os
 from argparse import Namespace
-from typing import Any, Dict
+from dataclasses import dataclass
+from typing import Any, ClassVar, Dict
 
 import modelopt.torch.distill as mtd
 import modelopt.torch.distill.plugins.megatron as mtd_mcore
-import modelopt.torch.opt as mto
 import yaml
 
 from megatron.core.models.gpt import GPTModel as MCoreGPTModel
@@ -18,15 +18,83 @@
     get_gpt_heterogeneous_layer_spec,
 )
 from megatron.core.models.hybrid.hybrid_model import HybridModel as MCoreHybridModel
+from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage
 from megatron.core.post_training.modelopt.gpt.model_specs import get_gpt_modelopt_spec
 from megatron.core.post_training.modelopt.gpt.state_dict_hooks import (
     mcore_gpt_load_te_state_dict_pre_hook,
 )
 from megatron.core.post_training.modelopt.hybrid.model_specs import get_hybrid_stack_modelopt_spec
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.module import MegatronModule
 from megatron.post_training.checkpointing import load_modelopt_state
-from megatron.post_training.utils import print_distributed_quant_summary
 from megatron.training import get_args, print_rank_0
 from megatron.training.arguments import core_transformer_config_from_args
+from megatron.training.models.gpt import GPTModelBuilder, GPTModelConfig
+from megatron.training.models.hybrid import HybridModelBuilder, HybridModelConfig
+
+
+@dataclass(kw_only=True)
+class ModelOptModelConfig(GPTModelConfig):
+    """Config for the legacy ModelOpt model construction path.
+
+    Identical to `GPTModelConfig` except for `builder` - construction still goes
+    through `gpt_config_from_args`, only the resolved builder class differs, since
+    ModelOpt-enabled runs need `ModelOptGPTModelBuilder` instead of `GPTModelBuilder`.
+    """
+
+    builder: ClassVar[str] = "megatron.post_training.model_builder.ModelOptGPTModelBuilder"
+
+
+@dataclass(kw_only=True)
+class ModelOptHybridModelConfig(HybridModelConfig):
+    """Config for the legacy ModelOpt model construction path, for hybrid models.
+
+    Identical to `HybridModelConfig` except for `builder` - construction still goes
+    through `hybrid_config_from_args`.
+    """
+
+    builder: ClassVar[str] = "megatron.post_training.model_builder.ModelOptHybridModelBuilder"
+
+
+class _ModelOptBuilderMixin:
+    """Shared `build_model()` override for the legacy ModelOpt model construction path.
+
+    `modelopt_gpt_hybrid_builder` dispatches on `args.export_model_type` internally, so
+    the same implementation covers both GPT and hybrid models - only the parent
+    `ModelBuilder` (and its `build_distributed_models()`) differs per config type, so
+    each gets its own concrete class below rather than sharing one tied to `GPTModelBuilder`.
+    """
+
+    def build_model(
+        self,
+        pg_collection: ProcessGroupCollection,
+        pre_process: bool | None = None,
+        post_process: bool | None = None,
+        vp_stage: int | None = None,
+    ) -> MegatronModule:
+        args = get_args()
+        if pre_process is None:
+            pre_process = is_pp_first_stage(pg_collection.pp)
+        if post_process is None:
+            post_process = is_pp_last_stage(pg_collection.pp)
+        return modelopt_gpt_hybrid_builder(
+            args,
+            pre_process,
+            post_process,
+            vp_stage,
+            pg_collection=pg_collection,
+        )
+
+
+class ModelOptGPTModelBuilder(_ModelOptBuilderMixin, GPTModelBuilder):
+    """ModelBuilder adapter for the legacy ModelOpt model construction path."""
+
+
+class ModelOptHybridModelBuilder(_ModelOptBuilderMixin, HybridModelBuilder):
+    """ModelBuilder adapter for the legacy ModelOpt model construction path (hybrid)."""
+
+
+logger = logging.getLogger(__name__)
 
 
 def count_parameters_in_layer(model, layer_name):
@@ -50,69 +118,40 @@ def _load_teacher_model_config(checkpoint_path: str) -> Namespace:
     """Reads teacher config from a file.
 
     The config provided, either in the teacher checkpoint dir or via `--export-kd-teacher-model-config`,
-    should specify (in NeMo yaml config format) any model architecture settings which differ from the main student model's.
-    This function will translate NeMo field names to MCore as needed.
+    should specify any model architecture settings which differ from the main student model's.
+    The field names should match those returned by get_args() and not TransformerConfig.
     """
-    required_teacher_fields = (
-        "num_layers",
-        "hidden_size",
-        "ffn_hidden_size",
-        "num_attention_heads",
-    )
-
     args = get_args()
+
     if args.export_kd_teacher_model_config is not None:
         config_path = args.export_kd_teacher_model_config
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(f"Teacher model-config file ({config_path}) not found.")
     else:
         config_path = os.path.join(checkpoint_path, "model_config.yaml")
-    if not os.path.exists(config_path):
-        raise FileNotFoundError(
-            f"Teacher model-config file {config_path} not found.\n"
-            "Teacher checkpoint dir must contain a NeMo-format config named 'model_config.yaml'"
-            " or provide it via --export-kd-teacher-model-config."
-        )
-    with open(config_path) as f:
-        config = yaml.safe_load(f)
-
-    if missing_keys := [k for k in required_teacher_fields if k not in config]:
-        raise ValueError(
-            f"Teacher model config file ({config_path}) missing the following required fields: {missing_keys}"
-        )
-
-    if "encoder_seq_length" in config:
-        config["seq_length"] = config["encoder_seq_length"]
-    if "bias" in config:
-        config["disable_bias_linear"] = not config["bias"]
-    if config.get("activation") == "swiglu":
-        config["swiglu"] = True
-    if config.get("position_embedding_type", False) is None:
-        config["use_rotary_position_embeddings"] = config["no_position_embedding"] = True
-    if "share_embeddings_and_output_weights" in config:
-        config["untie_embeddings_and_output_weights"] = not config[
-            "share_embeddings_and_output_weights"
-        ]
-    if "tokenizer" in config:
-        config["tokenizer_type"] = config["tokenizer"]["type"]
-        config["tokenizer_model"] = config["tokenizer"]["model"]
-    if "masked_softmax_fusion" in config:
-        config["no_masked_softmax_fusion"] = not config["masked_softmax_fusion"]
-    if config.get("normalization") == "layernorm1p":
-        config["apply_layernorm_1p"] = True
-    if "precision" in config:
-        config[config["precision"]] = True
-    if "mcore_gpt" in config:
-        config["use_mcore_models"] = config["mcore_gpt"]
-
-    args_dict = vars(get_args()).copy()
-    del args_dict["kv_channels"]  # not recalculated if present
-    # Setting teacher Flextron fields to false if training with Flextron, can be overridden
-    if "flextron" in args_dict:
-        config["flextron"] = False
-    if "enable_router" in args_dict:
-        config["enable_router"] = False
-    if "freeze_model" in args_dict:
-        config["freeze_model"] = False
-    args_dict.update(config)
+        if not os.path.exists(config_path):
+            logger.warning(
+                "No teacher config provided via --export-kd-teacher-model-config nor found at"
+                f" {checkpoint_path}/model_config.yaml. Assuming teacher model architecture same as student's."
+            )  # Useful for cases like QAD
+            config_path = None
+
+    args_dict = vars(args).copy()
+
+    if config_path is not None:
+        with open(config_path) as f:
+            config = yaml.safe_load(f)
+
+        del args_dict["kv_channels"]  # not recalculated if present
+        # Setting teacher Flextron fields to false if training with Flextron, can be overridden
+        if "flextron" in args_dict:
+            args_dict["flextron"] = False
+        if "enable_router" in args_dict:
+            args_dict["enable_router"] = False
+        if "freeze_model" in args_dict:
+            args_dict["freeze_model"] = False
+
+        args_dict.update(config)
 
     # Backward compat: old checkpoints have hybrid_override_pattern but not hybrid_layer_pattern
     if (args_dict.get('hybrid_override_pattern') is not None
@@ -152,7 +191,7 @@ def _build_teacher_model(config, config_raw: Namespace, model_kwargs: Dict[str,
 
     _add_load_convert_hooks(teacher)
 
-    # NOTE: Checkpoint loading now handled in `megatron/training/checkpointing.py`.
+    # NOTE: Checkpoint loading now handled by `megatron.post_training.checkpointing.load_kd_teacher_checkpoint()`.
 
     return teacher
 
@@ -366,7 +405,7 @@ def modelopt_gpt_hybrid_builder(
             )
 
         if args.export_default_te_spec and args.export_te_mcore_model:
-            logging.getLogger(__name__).warning(
+            logger.warning(
                 "--export-default-te-spec and --export-te-mcore-model are mutually exclusive. "
                 "Since --export-default-te-spec is given, --export-te-mcore-model will be disabled."
             )
@@ -466,10 +505,7 @@ def modelopt_gpt_hybrid_builder(
         # Additional tweaks needed for MCore.
         # (accounts for sharded state, pipeline parallel, and potentially skipping LM loss)
         mtd_mcore.adjust_distillation_model_for_mcore(model, distill_cfg)
-        # Also remove KD mode state to prevent issues with re-conversion after restore.
-        mto.ModeloptStateManager(model).state_dict().pop()  # TODO(aanoosheh): remove once fixed in ModelOpt
 
-    print_distributed_quant_summary(model)
     return model
 
 
diff --git a/megatron/post_training/utils.py b/megatron/post_training/utils.py
index 7bb5261522f..cbfb54b2cde 100644
--- a/megatron/post_training/utils.py
+++ b/megatron/post_training/utils.py
@@ -9,8 +9,25 @@
 from modelopt.torch.quantization.utils import is_quantized
 from packaging.version import Version
 
-from megatron.core import parallel_state
-from megatron.core.utils import unwrap_model
+
+def maybe_enable_modelopt(args):
+    """Set `args.modelopt_enabled` if a ModelOpt checkpoint or distillation teacher is
+    configured. Idempotent and safe to call multiple times (e.g. once early in
+    `pretrain_gpt.py` before building the model config, and again as a fallback in
+    `training.py` for callers that don't go through that entrypoint).
+    """
+    if getattr(args, "modelopt_enabled", False):
+        return
+
+    from megatron.post_training.checkpointing import has_modelopt_state
+    from megatron.training import print_rank_0
+
+    if args.load is not None and has_modelopt_state(args.load):
+        print_rank_0("ModelOpt checkpoint detected")
+        args.modelopt_enabled = True
+    if getattr(args, "export_kd_teacher_load", None):
+        # For distillation ckpts without ModelOpt state
+        args.modelopt_enabled = True
 
 
 def modelopt_version_higher_than(target_version: str):
diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py
index d8a757ddfc6..70f26c64d56 100644
--- a/megatron/training/argument_utils.py
+++ b/megatron/training/argument_utils.py
@@ -1,38 +1,41 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
+import ast
+import builtins
 import dataclasses
-import typing
-import types
-from typing import Any, Callable, Optional
-from argparse import ArgumentParser, _ArgumentGroup, Namespace
+import enum
 import inspect
 import itertools
-import builtins
-import ast
-import enum
-from dataclasses import Field, fields
+import types
+import typing
 import warnings
-import torch.nn.functional as F
+from argparse import ArgumentParser, Namespace, _ArgumentGroup
+from dataclasses import Field, fields
+from typing import Any, Callable, Optional
+
 import torch
+import torch.nn.functional as F
 
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.spec_utils import import_module
-
 from megatron.training.config import (
-    DistributedInitConfig, 
-    InferenceSetupConfig,
+    CheckpointConfig,
+    DistributedInitConfig,
     InferenceConfigContainer,
-    PretrainConfigContainer, 
-    SchedulerConfig, 
-    TokenizerConfig,
-    TrainingConfig, 
-    ValidationConfig, 
-    RNGConfig, 
+    InferenceSetupConfig,
     LoggerConfig,
+    PretrainConfigContainer,
+    ProfilingConfig,
+    RerunStateMachineConfig,
+    RNGConfig,
+    SchedulerConfig,
     StragglerDetectionConfig,
-    RerunStateMachineConfig, CheckpointConfig, ProfilingConfig
+    TokenizerConfig,
+    TrainingConfig,
+    ValidationConfig,
 )
-from megatron.training.models import HybridModelConfig, GPTModelConfig
+from megatron.training.models import GPTModelConfig, HybridModelConfig
+
 # TODO: support arg renames
 
 class TypeInferenceError(Exception):
@@ -274,14 +277,14 @@ def _get_field_docstrings(self, src_cfg_class: type) -> dict[str, str]:
 def core_transformer_config_from_args(args, config_class=None):
     from megatron.core.activations import squared_relu
     from megatron.core.fusions.fused_bias_geglu import quick_gelu
-    from megatron.core.transformer import MLATransformerConfig
-    from megatron.core.transformer.heterogeneous.heterogeneous_config import (
-        HeterogeneousTransformerConfig,
-    )
     from megatron.core.quantization.utils import (
         kitchen_quantization_recipe_config,
         load_quantization_recipe,
     )
+    from megatron.core.transformer import MLATransformerConfig
+    from megatron.core.transformer.heterogeneous.heterogeneous_config import (
+        HeterogeneousTransformerConfig,
+    )
 
     # Config class.
     config_class = config_class or TransformerConfig
@@ -427,8 +430,16 @@ def _default_config_from_args(cls: type, args: Namespace, return_instance: bool
         return kwargs
 
 
-def gpt_config_from_args(args: Namespace, config: TransformerConfig | None=None) -> Any:
-    """Create a GPTModelConfig from the appropriate values in the `args` Namespace."""
+def gpt_config_from_args(
+    args: Namespace, config: TransformerConfig | None = None, model_config_cls: type = GPTModelConfig
+) -> Any:
+    """Create a GPTModelConfig (or a compatible subclass) from the `args` Namespace.
+
+    `model_config_cls` lets callers reuse this same arg-derivation logic for
+    subclasses that only override metadata (e.g. `builder`) and add no new fields,
+    such as `ModelOptModelConfig`.
+    """
+    assert issubclass(model_config_cls, GPTModelConfig)
 
     kwargs = {}
     if config is None:
@@ -445,7 +456,6 @@ def gpt_config_from_args(args: Namespace, config: TransformerConfig | None=None)
     if args.spec is not None:
         kwargs["transformer_layer_spec"] = import_module(args.spec)
 
-
     kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy
     kwargs["position_embedding_type"] = args.position_embedding_type
     kwargs["rotary_percent"] = args.rotary_percent
@@ -468,11 +478,19 @@ def gpt_config_from_args(args: Namespace, config: TransformerConfig | None=None)
         kwargs["vocab_size"] = args.vocab_size
         kwargs["should_pad_vocab"] = True
 
-    return GPTModelConfig(**kwargs)
-    
+    return model_config_cls(**kwargs)
 
-def hybrid_config_from_args(args: Namespace, config: TransformerConfig | None=None) -> Any:
-    """Create a HybridModelConfig from the appropriate values in the `args` Namespace."""
+
+def hybrid_config_from_args(
+    args: Namespace, config: TransformerConfig | None = None, model_config_cls: type = HybridModelConfig
+) -> Any:
+    """Create a HybridModelConfig (or a compatible subclass) from the `args` Namespace.
+
+    `model_config_cls` lets callers reuse this same arg-derivation logic for
+    subclasses that only override metadata (e.g. `builder`) and add no new fields,
+    such as `ModelOptHybridModelConfig`.
+    """
+    assert issubclass(model_config_cls, HybridModelConfig)
 
     kwargs = {}
     if config is None:
@@ -488,7 +506,6 @@ def hybrid_config_from_args(args: Namespace, config: TransformerConfig | None=No
     elif args.spec is not None:
         kwargs["hybrid_stack_spec"] = import_module(args.spec)
 
-
     kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy
     kwargs["hybrid_layer_pattern"] = args.hybrid_layer_pattern
     kwargs["position_embedding_type"] = args.position_embedding_type
@@ -511,7 +528,7 @@ def hybrid_config_from_args(args: Namespace, config: TransformerConfig | None=No
         kwargs["vocab_size"] = args.vocab_size
         kwargs["should_pad_vocab"] = True
 
-    return HybridModelConfig(**kwargs)
+    return model_config_cls(**kwargs)
 
 
 def pretrain_cfg_container_from_args(args: Namespace, model_cfg=None) -> PretrainConfigContainer:
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index ad9de5b15b2..a136b277df5 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -2208,25 +2208,6 @@ def load_model_state_dict(module, state_dict, strict: bool):
     if has_nvidia_modelopt:
         print_distributed_quant_summary(model, msg="After loading checkpoint")
 
-        # Load teacher model in Distillation mode.
-        if getattr(args, "export_kd_teacher_load", None):
-            from megatron.post_training.checkpointing import load_modelopt_checkpoint
-
-            unwrapped_model = unwrap_model(model)[0]
-            # Note: load_modelopt_checkpoint may call this function so we prevent infinite recursion.
-            if hasattr(unwrapped_model, 'teacher_model'):
-                teacher = unwrapped_model.teacher_model
-                print_rank_0(f"Loading teacher as {type(teacher).__name__} from {args.export_kd_teacher_load} ...")
-                # [WAR]: To avoid error out on loading teacher's checkpoint, we temporarily
-                # set args.finetune to True while loading the teacher checkpoint.
-                original_args_finetune, original_ckpt_format = args.finetune, args.ckpt_format
-                args.finetune = True
-                if args.export_kd_teacher_ckpt_format is not None:
-                    args.ckpt_format = args.export_kd_teacher_ckpt_format
-                load_modelopt_checkpoint([teacher], load_arg='export_kd_teacher_load')
-                args.finetune, args.ckpt_format = original_args_finetune, original_ckpt_format
-                print_rank_0("... teacher loaded successfully.")
-
     return iteration, num_floating_point_operations_so_far
 
 
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 730a2ba329b..820a075e919 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -207,6 +207,8 @@
 try:
     from modelopt.torch.distill.plugins.megatron import get_tensor_shapes_adjust_fn_for_distillation
 
+    from megatron.post_training.utils import maybe_enable_modelopt
+
     has_nvidia_modelopt = True
 except ImportError:
     has_nvidia_modelopt = False
@@ -1701,16 +1703,7 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap
                 print_rank_0(">   including expert parallelism AG group")
 
     if has_nvidia_modelopt:
-        from megatron.post_training.checkpointing import has_modelopt_state
-
-        # [ModelOpt]: Check if the checkpoint is a ModelOpt checkpoint and
-        # set a flag to use our model provider if so.
-        if args.load is not None and has_modelopt_state(args.load):
-            print_rank_0(f'ModelOpt checkpoint detected')
-            args.modelopt_enabled = True
-        elif getattr(args, "export_kd_teacher_load", None):
-            # For distillation ckpts without ModelOpt state
-            args.modelopt_enabled = True
+        maybe_enable_modelopt(args)
 
     # Build model.
     def build_model():
@@ -2017,6 +2010,9 @@ def setup_model_and_optimizer(
     skip_optimizer = not (has_normal_optimizer or has_rl_optimizer)
     wrap_with_ddp = not skip_optimizer
 
+    if has_nvidia_modelopt:
+        maybe_enable_modelopt(args)
+
     def _build_model_wrapper(wrap_with_ddp: bool):
         if cfg_container is not None and getattr(cfg_container, "model", None) is not None:
             from megatron.training.utils import start_memory_history_recording
@@ -2177,6 +2173,14 @@ def _build_model_wrapper(wrap_with_ddp: bool):
         args.iteration = 0
         args.num_floating_point_operations_so_far = 0
 
+    # [ModelOpt]: Load the teacher checkpoint for ModelOpt distillation if applicable.
+    # Import locally to prevent circular import: megatron.post_training.checkpointing
+    # imports `get_args` from megatron.training at module scope.
+    if has_nvidia_modelopt:
+        from megatron.post_training.checkpointing import load_kd_teacher_checkpoint
+
+        load_kd_teacher_checkpoint(model)
+
     # Validate that the world size can accommodate the current batch size.
     # This catches the case where GPUs were scaled up mid-training but the
     # current position in the batch size schedule yields a batch size that
diff --git a/pretrain_gpt.py b/pretrain_gpt.py
index 0a87db5cdb1..d903bd7d96d 100644
--- a/pretrain_gpt.py
+++ b/pretrain_gpt.py
@@ -66,6 +66,8 @@
 try:
     from megatron.post_training.arguments import add_modelopt_args
     from megatron.post_training.loss_func import loss_func as loss_func_modelopt
+    from megatron.post_training.model_builder import ModelOptModelConfig
+    from megatron.post_training.utils import maybe_enable_modelopt
 
     has_nvidia_modelopt = True
 except ImportError:
@@ -503,7 +505,12 @@ def get_embedding_ranks(pp_ranks: List[int]):
         extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None,
         args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
     )
-    model_cfg = gpt_config_from_args(args)
+    if has_nvidia_modelopt:
+        maybe_enable_modelopt(args)
+    if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False):
+        model_cfg = gpt_config_from_args(args, model_config_cls=ModelOptModelConfig)
+    else:
+        model_cfg = gpt_config_from_args(args)
     full_config = pretrain_cfg_container_from_args(args, model_cfg)
     pretrain(
         full_config,
diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py
index 6cd65b4e47c..d59eef69480 100644
--- a/pretrain_hybrid.py
+++ b/pretrain_hybrid.py
@@ -66,6 +66,8 @@
 try:
     from megatron.post_training.arguments import add_modelopt_args
     from megatron.post_training.loss_func import loss_func as loss_func_modelopt
+    from megatron.post_training.model_builder import ModelOptHybridModelConfig
+    from megatron.post_training.utils import maybe_enable_modelopt
 
     has_nvidia_modelopt = True
 except ImportError:
@@ -449,7 +451,12 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None
         extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None,
         args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
     )
-    model_cfg = hybrid_config_from_args(args)
+    if has_nvidia_modelopt:
+        maybe_enable_modelopt(args)
+    if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False):
+        model_cfg = hybrid_config_from_args(args, model_config_cls=ModelOptHybridModelConfig)
+    else:
+        model_cfg = hybrid_config_from_args(args)
     full_config = pretrain_cfg_container_from_args(args, model_cfg)
     pretrain(
         full_config,
diff --git a/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py b/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py
index 091623b9b84..fcee30e61e6 100644
--- a/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py
+++ b/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py
@@ -50,6 +50,7 @@ def collect_train_test_metrics(
             "lm loss",
             "num-zeros",
             "mtp_1 loss",
+            "total loss",
         ]
     }
 
diff --git a/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py b/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py
index 68aa0db5622..a5ad326f49d 100644
--- a/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py
+++ b/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py
@@ -20,6 +20,7 @@
     "num-zeros": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.20)],
     "generated_tokens": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)],
     "logprobs": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)],
+    "total loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)],
 }
 
 
diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/golden_values_dev_dgx_a100.json b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/golden_values_dev_dgx_a100.json
deleted file mode 100644
index ae531bf007e..00000000000
--- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/golden_values_dev_dgx_a100.json
+++ /dev/null
@@ -1,109 +0,0 @@
-{
-    "kd loss": {
-        "start_step": 1,
-        "end_step": 100,
-        "step_interval": 1,
-        "values": {
-            "1": 0.4930937,
-            "2": 0.4935429,
-            "3": 0.4938237,
-            "4": 0.4938229,
-            "5": 0.4924306,
-            "6": 0.4935288,
-            "7": 0.4928354,
-            "8": 0.4925097,
-            "9": 0.4936136,
-            "10": 0.4922911,
-            "11": 0.4934031,
-            "12": 0.4951033,
-            "13": 0.4918853,
-            "14": 0.4936183,
-            "15": 0.4926639,
-            "16": 0.4927304,
-            "17": 0.4925308,
-            "18": 0.4927951,
-            "19": 0.4938825,
-            "20": 0.4939776,
-            "21": 0.4933512,
-            "22": 0.4935322,
-            "23": 0.4937269,
-            "24": 0.4927326,
-            "25": 0.4927868,
-            "26": 0.4927689,
-            "27": 0.4924214,
-            "28": 0.4925573,
-            "29": 0.4917694,
-            "30": 0.4919884,
-            "31": 0.4929765,
-            "32": 0.4930308,
-            "33": 0.4928029,
-            "34": 0.4923102,
-            "35": 0.4918847,
-            "36": 0.4914086,
-            "37": 0.4929215,
-            "38": 0.4923307,
-            "39": 0.4910690,
-            "40": 0.4919418,
-            "41": 0.4913271,
-            "42": 0.4919568,
-            "43": 0.4903573,
-            "44": 0.4916522,
-            "45": 0.4915655,
-            "46": 0.4898856,
-            "47": 0.4899229,
-            "48": 0.4892673,
-            "49": 0.4894423,
-            "50": 0.4903796,
-            "51": 0.4907262,
-            "52": 0.4882944,
-            "53": 0.4877340,
-            "54": 0.4902404,
-            "55": 0.4881638,
-            "56": 0.4888564,
-            "57": 0.4882180,
-            "58": 0.4887677,
-            "59": 0.4883497,
-            "60": 0.4863744,
-            "61": 0.4875762,
-            "62": 0.4837778,
-            "63": 0.4867221,
-            "64": 0.4840697,
-            "65": 0.4840384,
-            "66": 0.4857976,
-            "67": 0.4837634,
-            "68": 0.4800620,
-            "69": 0.4781690,
-            "70": 0.4818793,
-            "71": 0.4796092,
-            "72": 0.4783594,
-            "73": 0.4789546,
-            "74": 0.4767389,
-            "75": 0.4774750,
-            "76": 0.4746155,
-            "77": 0.4745574,
-            "78": 0.4737080,
-            "79": 0.4718909,
-            "80": 0.4693059,
-            "81": 0.4696763,
-            "82": 0.4705839,
-            "83": 0.4661151,
-            "84": 0.4634311,
-            "85": 0.4654806,
-            "86": 0.4634389,
-            "87": 0.4595202,
-            "88": 0.4586285,
-            "89": 0.4595589,
-            "90": 0.4561397,
-            "91": 0.4547178,
-            "92": 0.4553441,
-            "93": 0.4545716,
-            "94": 0.4518545,
-            "95": 0.4504384,
-            "96": 0.4518549,
-            "97": 0.4450935,
-            "98": 0.4441919,
-            "99": 0.4442465,
-            "100": 0.4427010
-        }
-    }
-}
diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/golden_values_dev_dgx_h100.json
new file mode 100644
index 00000000000..8634a66847a
--- /dev/null
+++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/golden_values_dev_dgx_h100.json
@@ -0,0 +1,216 @@
+{
+    "total loss": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": 0.49197,
+            "2": 0.49244,
+            "3": 0.49244,
+            "4": 0.49309,
+            "5": 0.49199,
+            "6": 0.4925,
+            "7": 0.49298,
+            "8": 0.49175,
+            "9": 0.4929,
+            "10": 0.49224,
+            "11": 0.49354,
+            "12": 0.49217,
+            "13": 0.4921,
+            "14": 0.49196,
+            "15": 0.49207,
+            "16": 0.49239,
+            "17": 0.49231,
+            "18": 0.49241,
+            "19": 0.49228,
+            "20": 0.49096,
+            "21": 0.49281,
+            "22": 0.49273,
+            "23": 0.49169,
+            "24": 0.49298,
+            "25": 0.49222,
+            "26": 0.49219,
+            "27": 0.49351,
+            "28": 0.4928,
+            "29": 0.49313,
+            "30": 0.49276,
+            "31": 0.49254,
+            "32": 0.49177,
+            "33": 0.49254,
+            "34": 0.49255,
+            "35": 0.49246,
+            "36": 0.49135,
+            "37": 0.49143,
+            "38": 0.49129,
+            "39": 0.49205,
+            "40": 0.49215,
+            "41": 0.49167,
+            "42": 0.49202,
+            "43": 0.49159,
+            "44": 0.49097,
+            "45": 0.49001,
+            "46": 0.49096,
+            "47": 0.49014,
+            "48": 0.4894,
+            "49": 0.48931,
+            "50": 0.48995,
+            "51": 0.49027,
+            "52": 0.48921,
+            "53": 0.4916,
+            "54": 0.49015,
+            "55": 0.4892,
+            "56": 0.48764,
+            "57": 0.48865,
+            "58": 0.4877,
+            "59": 0.4865,
+            "60": 0.48435,
+            "61": 0.48678,
+            "62": 0.48624,
+            "63": 0.48259,
+            "64": 0.48274,
+            "65": 0.48095,
+            "66": 0.48127,
+            "67": 0.48124,
+            "68": 0.47975,
+            "69": 0.47882,
+            "70": 0.47826,
+            "71": 0.47797,
+            "72": 0.47728,
+            "73": 0.47533,
+            "74": 0.47329,
+            "75": 0.47452,
+            "76": 0.4729,
+            "77": 0.47196,
+            "78": 0.46773,
+            "79": 0.46857,
+            "80": 0.46752,
+            "81": 0.46539,
+            "82": 0.46683,
+            "83": 0.46365,
+            "84": 0.45854,
+            "85": 0.45937,
+            "86": 0.46228,
+            "87": 0.45535,
+            "88": 0.45485,
+            "89": 0.45797,
+            "90": 0.44956,
+            "91": 0.45188,
+            "92": 0.44878,
+            "93": 0.4514,
+            "94": 0.44644,
+            "95": 0.44778,
+            "96": 0.44672,
+            "97": 0.44107,
+            "98": 0.44625,
+            "99": 0.43963,
+            "100": 0.43588
+        }
+    },
+    "lm loss": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": 0.0,
+            "2": 0.0,
+            "3": 0.0,
+            "4": 0.0,
+            "5": 0.0,
+            "6": 0.0,
+            "7": 0.0,
+            "8": 0.0,
+            "9": 0.0,
+            "10": 0.0,
+            "11": 0.0,
+            "12": 0.0,
+            "13": 0.0,
+            "14": 0.0,
+            "15": 0.0,
+            "16": 0.0,
+            "17": 0.0,
+            "18": 0.0,
+            "19": 0.0,
+            "20": 0.0,
+            "21": 0.0,
+            "22": 0.0,
+            "23": 0.0,
+            "24": 0.0,
+            "25": 0.0,
+            "26": 0.0,
+            "27": 0.0,
+            "28": 0.0,
+            "29": 0.0,
+            "30": 0.0,
+            "31": 0.0,
+            "32": 0.0,
+            "33": 0.0,
+            "34": 0.0,
+            "35": 0.0,
+            "36": 0.0,
+            "37": 0.0,
+            "38": 0.0,
+            "39": 0.0,
+            "40": 0.0,
+            "41": 0.0,
+            "42": 0.0,
+            "43": 0.0,
+            "44": 0.0,
+            "45": 0.0,
+            "46": 0.0,
+            "47": 0.0,
+            "48": 0.0,
+            "49": 0.0,
+            "50": 0.0,
+            "51": 0.0,
+            "52": 0.0,
+            "53": 0.0,
+            "54": 0.0,
+            "55": 0.0,
+            "56": 0.0,
+            "57": 0.0,
+            "58": 0.0,
+            "59": 0.0,
+            "60": 0.0,
+            "61": 0.0,
+            "62": 0.0,
+            "63": 0.0,
+            "64": 0.0,
+            "65": 0.0,
+            "66": 0.0,
+            "67": 0.0,
+            "68": 0.0,
+            "69": 0.0,
+            "70": 0.0,
+            "71": 0.0,
+            "72": 0.0,
+            "73": 0.0,
+            "74": 0.0,
+            "75": 0.0,
+            "76": 0.0,
+            "77": 0.0,
+            "78": 0.0,
+            "79": 0.0,
+            "80": 0.0,
+            "81": 0.0,
+            "82": 0.0,
+            "83": 0.0,
+            "84": 0.0,
+            "85": 0.0,
+            "86": 0.0,
+            "87": 0.0,
+            "88": 0.0,
+            "89": 0.0,
+            "90": 0.0,
+            "91": 0.0,
+            "92": 0.0,
+            "93": 0.0,
+            "94": 0.0,
+            "95": 0.0,
+            "96": 0.0,
+            "97": 0.0,
+            "98": 0.0,
+            "99": 0.0,
+            "100": 0.0
+        }
+    }
+}
diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/golden_values_lts_dgx_a100.json b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/golden_values_lts_dgx_a100.json
deleted file mode 100644
index 9e26dfeeb6e..00000000000
--- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/golden_values_lts_dgx_a100.json
+++ /dev/null
@@ -1 +0,0 @@
-{}
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/model_config.yaml
index c75a5a81414..2ba8cb329d0 100644
--- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_modelopt_distill_resume/model_config.yaml
@@ -1,11 +1,10 @@
 ENV_VARS:
-  SKIP_PYTEST: 1
   CUDA_DEVICE_MAX_CONNECTIONS: 1
   NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
   NCCL_ALGO: Ring
   CUBLAS_WORKSPACE_CONFIG: :4096:8
   ARTIFACTS_ROOT: /workspace/checkpoints
-  DISTILL_CONFIG: '{intermediate_layer_pairs: [["decoder.final_layernorm", "decoder.final_layernorm"]], logit_layers: ["output_layer", "output_layer"], skip_lm_loss: true, kd_loss_scale: 10.0}'
+  DISTILL_CONFIG: '{intermediate_layer_pairs: [["decoder.final_layernorm", "decoder.final_layernorm"]], logit_layers: ["output_layer", "output_layer"], skip_lm_loss: true, kd_loss_scale: 1.0}'
 BEFORE_SCRIPT: |
   mkdir -p ${DATA_CACHE_PATH}/distill && echo $DISTILL_CONFIG | yq -P > ${DATA_CACHE_PATH}/distill/distill_config.yaml
 MODEL_ARGS:
@@ -68,4 +67,8 @@ MODEL_ARGS:
   --log-memory-to-tensorboard: true
   --async-save: true
   --use-persistent-ckpt-worker: true
+  --exit-interval: 100
 TEST_TYPE: ckpt-resume
+METRICS:
+  - lm loss
+  - total loss
diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_nemotron_v3_pico_7b_a1b_tp1_ep8_QAD_dgx_h100_1N8G/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/hybrid/hybrid_nemotron_v3_pico_7b_a1b_tp1_ep8_QAD_dgx_h100_1N8G/golden_values_dev_dgx_h100.json
new file mode 100644
index 00000000000..2b2029d40f9
--- /dev/null
+++ b/tests/functional_tests/test_cases/hybrid/hybrid_nemotron_v3_pico_7b_a1b_tp1_ep8_QAD_dgx_h100_1N8G/golden_values_dev_dgx_h100.json
@@ -0,0 +1,36 @@
+{
+    "total loss": {
+        "start_step": 1,
+        "end_step": 10,
+        "step_interval": 1,
+        "values": {
+            "1": 1.58285,
+            "2": 1.6902,
+            "3": 0.1002,
+            "4": 0.06419,
+            "5": 0.06295,
+            "6": 0.06725,
+            "7": 0.16463,
+            "8": 0.06043,
+            "9": 0.06715,
+            "10": 0.05355
+        }
+    },
+    "lm loss": {
+        "start_step": 1,
+        "end_step": 10,
+        "step_interval": 1,
+        "values": {
+            "1": 0.0,
+            "2": 0.0,
+            "3": 0.0,
+            "4": 0.0,
+            "5": 0.0,
+            "6": 0.0,
+            "7": 0.0,
+            "8": 0.0,
+            "9": 0.0,
+            "10": 0.0
+        }
+    }
+}
diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_nemotron_v3_pico_7b_a1b_tp1_ep8_QAD_dgx_h100_1N8G/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_nemotron_v3_pico_7b_a1b_tp1_ep8_QAD_dgx_h100_1N8G/model_config.yaml
new file mode 100644
index 00000000000..e2116268b4c
--- /dev/null
+++ b/tests/functional_tests/test_cases/hybrid/hybrid_nemotron_v3_pico_7b_a1b_tp1_ep8_QAD_dgx_h100_1N8G/model_config.yaml
@@ -0,0 +1,190 @@
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: ":4096:8"
+  TRITON_CACHE_AUTOTUNING: 0
+  MAMBA_DETERMINISTIC: 1
+  # Paths
+  MODEL_BF16_CKPT: "${DATA_PATH}/model/nemotron_v3_pico_7b-a1b/3T-token_deeparch"
+  PTQ_QUANTIZED_CKPT: "${DATA_CACHE_PATH}/ptq_quantized_ckpt"
+  TOKENIZER: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
+  HF_HOME: "${DATA_PATH}/hf_home"
+BEFORE_SCRIPT: |
+  # Stage 1: PTQ quantization via quantize.py directly
+  echo -e "\n=== Stage 1: Running PTQ (NVFP4) ===\n"
+  cd /opt/megatron-lm
+  # Env vars that arguments.sh would normally set
+  export TOKENIZERS_PARALLELISM=False
+  export OMP_NUM_THREADS=1
+  export NCCL_IB_SL=1
+  export NCCL_IB_TIMEOUT=22
+  uv run --no-sync python -m torch.distributed.run --nproc_per_node=8 \
+    examples/post_training/modelopt/quantize.py \
+    --deterministic-mode \
+    --micro-batch-size 1 \
+    --save-interval 100000 \
+    --bf16 \
+    --seq-length 4096 \
+    --max-position-embeddings 4096 \
+    --tokenizer-type HuggingFaceTokenizer \
+    --tokenizer-model ${TOKENIZER} \
+    --tensor-model-parallel-size 1 \
+    --expert-model-parallel-size 8 \
+    --expert-tensor-parallel-size 1 \
+    --pipeline-model-parallel-size 1 \
+    --context-parallel-size 1 \
+    --hidden-size 1216 \
+    --num-attention-heads 32 \
+    --group-query-attention \
+    --num-query-groups 2 \
+    --ffn-hidden-size 896 \
+    --kv-channels 128 \
+    --squared-relu \
+    --normalization RMSNorm \
+    --disable-bias-linear \
+    --attention-dropout 0.0 \
+    --hidden-dropout 0.0 \
+    --position-embedding-type none \
+    --untie-embeddings-and-output-weights \
+    --init-method-std 0.0256 \
+    --hybrid-layer-pattern MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME \
+    --mamba-num-heads 64 \
+    --export-model-type MambaModel \
+    --num-experts 128 \
+    --moe-router-topk 6 \
+    --moe-aux-loss-coeff 1e-4 \
+    --moe-router-topk-scaling-factor 2.5 \
+    --moe-router-enable-expert-bias \
+    --moe-router-dtype fp32 \
+    --moe-router-score-function sigmoid \
+    --moe-router-load-balancing-type seq_aux_loss \
+    --moe-shared-expert-intermediate-size 3712 \
+    --moe-token-dispatcher-type alltoall \
+    --moe-grouped-gemm \
+    --use-fused-weighted-squared-relu \
+    --attention-backend fused \
+    --disable-gloo-process-groups \
+    --no-create-attention-mask-in-dataloader \
+    --ckpt-format torch_dist \
+    --ckpt-fully-parallel-load \
+    --load ${MODEL_BF16_CKPT} \
+    --save ${PTQ_QUANTIZED_CKPT} \
+    --finetune \
+    --auto-detect-ckpt-format \
+    --distributed-timeout-minutes 30 \
+    --export-quant-cfg MAMBA_MOE_NVFP4_CONSERVATIVE_CFG \
+    --calib-dataset-path-or-name cnn_dailymail \
+    --calib-size 256 \
+    --calib-batch-size 8 \
+    --skip-generate \
+    --export-te-mcore-model
+  #--export-default-te-spec         # TODO(aanoosheh): undo once TE fix is released
+  echo -e "\n=== Stage 1 complete ===\n"
+MODEL_ARGS:
+  # KD teacher/config
+  --export-te-mcore-model: true
+  #--export-default-te-spec: true   # TODO(aanoosheh): undo once TE fix is released
+  --export-kd-teacher-load: ${MODEL_BF16_CKPT}
+  --auto-detect-ckpt-format: true
+  --finetune: true
+  # Architecture
+  --hidden-size: 1216
+  --num-attention-heads: 32
+  --group-query-attention: true
+  --num-query-groups: 2
+  --ffn-hidden-size: 896
+  --kv-channels: 128
+  --squared-relu: true
+  --normalization: RMSNorm
+  --disable-bias-linear: true
+  --attention-dropout: 0.0
+  --hidden-dropout: 0.0
+  --position-embedding-type: none
+  --untie-embeddings-and-output-weights: true
+  --init-method-std: 0.0256
+  --hybrid-layer-pattern: MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME
+  --mamba-num-heads: 64
+  --export-model-type: MambaModel
+  # MoE
+  --num-experts: 128
+  --moe-router-topk: 6
+  --moe-aux-loss-coeff: 1e-4
+  --moe-router-topk-scaling-factor: 2.5
+  --moe-router-enable-expert-bias: true
+  --moe-router-dtype: fp32
+  --moe-router-score-function: sigmoid
+  --moe-router-load-balancing-type: seq_aux_loss
+  --moe-shared-expert-intermediate-size: 3712
+  --moe-token-dispatcher-type: alltoall
+  --moe-grouped-gemm: true
+  --use-fused-weighted-squared-relu: true
+  # Tokenizer
+  --tokenizer-type: SFTTokenizer
+  --tokenizer-model: ${TOKENIZER}
+  --sft: true
+  --sft-tokenizer-prompt-format: identity
+  --bf16: true
+  # Parallelism
+  --tensor-model-parallel-size: 1     # TODO(aanoosheh): can change to 2 once TE fix is released
+  --expert-model-parallel-size: 4
+  --expert-tensor-parallel-size: 1
+  --pipeline-model-parallel-size: 1
+  --context-parallel-size: 2
+  --sequence-parallel: true
+  # Infrastructure
+  --attention-backend: fused
+  --disable-gloo-process-groups: true
+  --no-create-attention-mask-in-dataloader: true
+  --ddp-num-buckets: 8
+  --override-opt_param-scheduler: true
+  --num-workers: 1
+  --ckpt-format: torch_dist
+  --ckpt-fully-parallel-save: true
+  --ckpt-fully-parallel-load: true
+  --ckpt-assume-constant-structure: true
+  # Training
+  --micro-batch-size: 1
+  --global-batch-size: 8
+  --seq-length: 2048               # TODO(aanoosheh): change to 4096 once TE fix is released
+  --max-position-embeddings: 2048  # TODO(aanoosheh): change to 4096 once TE fix is released
+  --train-iters: 10
+  --lr: 0.00015
+  --lr-decay-style: cosine
+  --lr-decay-iters: 320000
+  --min-lr: 1.0e-5
+  --weight-decay: 1e-2
+  --clip-grad: 1.0
+  --lr-warmup-fraction: .01
+  --use-distributed-optimizer: true
+  --overlap-param-gather: true
+  --overlap-grad-reduce: true
+  # Checkpoint & data
+  --save: ${CHECKPOINT_SAVE_PATH}
+  --load: ${PTQ_QUANTIZED_CKPT}
+  --data-path: "${DATA_PATH}/text/nemotron-3-super-sft_train-sample.jsonl"
+  --split: "949,50,1"
+  --distributed-backend: nccl
+  --transformer-impl: transformer_engine
+  --data-cache-path: ${DATA_CACHE_PATH}
+  # Logging
+  --log-interval: 1
+  --save-interval: 10
+  --eval-interval: 10
+  --eval-iters: 2
+  --log-params-norm: true
+  --log-num-zeros-in-grad: true
+  --log-validation-ppl-to-tensorboard: true
+  --log-timers-to-tensorboard: true
+  --tensorboard-dir: ${TENSORBOARD_PATH}
+  --log-memory-to-tensorboard: true
+  --timing-log-level: 0
+  --no-gradient-accumulation-fusion: true
+  --distributed-timeout-minutes: 30
+  # Etc
+  --deterministic-mode: true
+  --exit-interval: 10
+TEST_TYPE: regular
+METRICS:
+  - lm loss
+  - total loss
diff --git a/tests/test_utils/recipes/h100/mamba.yaml b/tests/test_utils/recipes/h100/mamba.yaml
index 0ec42cc8a22..64ca1beef56 100644
--- a/tests/test_utils/recipes/h100/mamba.yaml
+++ b/tests/test_utils/recipes/h100/mamba.yaml
@@ -96,3 +96,9 @@ products:
         platforms: [dgx_h100]
       # - environment: [lts] # disabled until triton is bumped
       #   scope: [nightly]
+
+  - test_case: [hybrid_nemotron_v3_pico_7b_a1b_tp1_ep8_QAD_dgx_h100_1N8G]
+    products:
+      - environment: [dev]
+        scope: [mr, mr-github]
+        platforms: [dgx_h100]
diff --git a/tools/common_pile_dataset/README.md b/tools/common_pile_dataset/README.md
index 2431d1b01d3..dcc18fee24b 100644
--- a/tools/common_pile_dataset/README.md
+++ b/tools/common_pile_dataset/README.md
@@ -165,7 +165,7 @@ default. On HPC systems where `/home` is small, set `HF_HOME` to a path with
 sufficient space:
 
 ```bash
-export HF_HOME=/lustre/path/to/.hf_cache
+export HF_HOME=/lustre/path/to/hf_home
 ```
 
 The setup script does this automatically.
diff --git a/tools/common_pile_dataset/setup_common_pile_dataset.sh b/tools/common_pile_dataset/setup_common_pile_dataset.sh
index cb869e28368..01c438cd21a 100644
--- a/tools/common_pile_dataset/setup_common_pile_dataset.sh
+++ b/tools/common_pile_dataset/setup_common_pile_dataset.sh
@@ -29,7 +29,7 @@ DATASET_NAME="common-pile/comma_v0.1_training_dataset"
 WORK_DIR="/tmp/mcore_dataset_setup_$$"
 
 # Redirect HuggingFace cache to lustre so it doesn't fill up /home
-export HF_HOME="/lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_mcore/mcore_ci/.hf_cache"
+export HF_HOME="/lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_mcore/mcore_ci/hf_home"
 export HF_DATASETS_CACHE="${HF_HOME}/datasets"
 
 echo "============================================================"