diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev
index 179c2aa21a6..9ebda435f99 100644
--- a/docker/Dockerfile.ci.dev
+++ b/docker/Dockerfile.ci.dev
@@ -72,6 +72,7 @@ EOF
 
 COPY assets/ /opt/data/
 ENV UV_PYTHON=$UV_PROJECT_ENVIRONMENT/bin/python
+COPY . /opt/megatron-lm/
 
 ##### For NVIDIANS only #####
 FROM main as jet
diff --git a/tests/functional_tests/test_cases/bert/bert_release/model_config.yaml b/tests/functional_tests/test_cases/bert/bert_release/model_config.yaml
index ab5558fa7d2..278ad6c17a8 100644
--- a/tests/functional_tests/test_cases/bert/bert_release/model_config.yaml
+++ b/tests/functional_tests/test_cases/bert/bert_release/model_config.yaml
@@ -27,7 +27,7 @@ MODEL_ARGS:
   --pipeline-model-parallel-size: 8
   # Data args
   --data-path: ${DATA_BLEND}
-  --vocab-file: ${DATA_PATH}/text/the_pile/bert_shard00/vocab.txt
+  --vocab-file: ${DATA_PATH}/vocab.txt
   --split: 949,50,1
   --data-cache-path: ${DATA_CACHE_PATH}
   # EVAL_AND_LOGGING_ARGS
diff --git a/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release/model_config.yaml b/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release/model_config.yaml
index c16fedc7860..080f669e6a4 100644
--- a/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release/model_config.yaml
+++ b/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release/model_config.yaml
@@ -12,11 +12,12 @@ ENV_VARS:
   NCCL_DEBUG: VERSION
   NON_DETERMINSTIC_RESULTS: 1
   NVSHMEM_IB_ENABLE_IBGDA: 0
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
 TEST_TYPE: "release"
 MODEL_ARGS:
   # Distributed args
   --distributed-timeout-minutes: 60
-  --tensor-model-parallel-size: 1
+  --tensor-model-parallel-size: 2
   --pipeline-model-parallel-size: 4
   --pipeline-model-parallel-layout: Et*2\\|\\(tt\\|\\)*5t\\|tmL # Et*2|(tt|)*5t|tmL
   --expert-model-parallel-size: 16
@@ -47,8 +48,8 @@ MODEL_ARGS:
   # Data args
   --seq-length: 4096
   --data-cache-path: ${DATA_CACHE_PATH}
-  --tokenizer-type: HuggingFaceTokenizer
-  --tokenizer-model: ${TOKENIZER_PATH}
+  --tokenizer-type: GPTSentencePieceTokenizer
+  --tokenizer-model: ${DATA_PATH}/utils/nemotron_2_256k.model
   --data-path: $DATA_BLEND
   --split: 99,1,0
   --no-mmap-bin-files: true
@@ -81,12 +82,11 @@ MODEL_ARGS:
   --qk-layernorm: true
 
   # Add learning rate args
-  --lr-decay-samples: 584765624
+  --lr-decay-samples: 24413696
   --lr-warmup-samples: 1536000
-  # Learning rate scaled down from 7.3e-6 (DeepSeek-V3 technical report, GBS=15360) to 3.9e-6 (GBS=8192)
-  --lr-warmup-init: 3.9e-7
-  --lr: 3.9e-6
-  --min-lr: 3.9e-7
+  --lr-warmup-init: 1e-7
+  --lr: 1e-5
+  --min-lr: 1e-6
   --lr-decay-style: cosine
   --adam-beta1: 0.9
   --adam-beta2: 0.95
@@ -127,8 +127,6 @@ MODEL_ARGS:
   --eval-interval: 200
 
   # Add checkpointing args
-  --no-load-optim: true
-  --no-load-rng: true
   --auto-detect-ckpt-format:
     true
     # Add checkpointing args
diff --git a/tests/functional_tests/test_cases/t5/t5_release/model_config.yaml b/tests/functional_tests/test_cases/t5/t5_release/model_config.yaml
index 852fbf9819d..b684a2ebb54 100644
--- a/tests/functional_tests/test_cases/t5/t5_release/model_config.yaml
+++ b/tests/functional_tests/test_cases/t5/t5_release/model_config.yaml
@@ -37,7 +37,7 @@ MODEL_ARGS:
   --pipeline-model-parallel-size: 1
   # Data args
   --data-path: ${DATA_BLEND}
-  --vocab-file: ${DATA_PATH}/text/the_pile/t5_shard00/bert-large-cased-vocab.txt
+  --vocab-file: ${DATA_PATH}/bert-large-cased-vocab.txt
   --tokenizer-type: BertWordPieceCase
   --split: 99982,9,9
   --data-cache-path: ${DATA_CACHE_PATH}
diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
index 57ff1df6de2..80b0d2bca69 100644
--- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py
+++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import copy
 import dataclasses
@@ -115,6 +115,7 @@ def __del__(self):
         torch.cuda.synchronize()
         Utils.destroy_model_parallel()
 
+    @pytest.mark.flaky_in_dev
     @pytest.mark.internal
     def dispatcher_dropless_test(self):
         moe_layer = self.moe_layer
@@ -157,6 +158,7 @@ def dispatcher_dropless_test(self):
             hidden_states.grad, ans
         ), "Restored hidden states do not match original hidden states"
 
+    @pytest.mark.flaky_in_dev
     @pytest.mark.internal
     def dispatcher_capacity_test(self):
         moe_layer = self.moe_layer
@@ -210,6 +212,7 @@ def dispatcher_capacity_test(self):
             hidden_states.grad, restored_hidden_states_answer
         ), "Gradient of hidden states should be same as hidden states"
 
+    @pytest.mark.flaky_in_dev
     @pytest.mark.internal
     def dispatcher_drop_and_pad_test(self):
         """Test if the tokens are dropped and padded correctly.
@@ -275,6 +278,7 @@ def dispatcher_drop_and_pad_test(self):
             hidden_states.grad, backward_answer
         ), "Gradient of hidden states should be same as hidden states"
 
+    @pytest.mark.flaky_in_dev
     @pytest.mark.internal
     def dispatcher_router_padding_for_fp8_test(self):
         """Test if the routing map is padded correctly for FP8 training.
@@ -367,6 +371,7 @@ def teardown_method(self, method):
     @pytest.mark.internal
     @pytest.mark.parametrize("tp_size,ep_size", [(8, 1), (1, 8), (2, 4), (1, 1)])
     @pytest.mark.parametrize("permute_fusion", permute_fusion_params)
+    @pytest.mark.flaky_in_dev
     def test_forward_backward(self, tp_size, ep_size, permute_fusion):
         container = MoEModelTestContainer(
             tp_size=tp_size,