update tests

Author: liding-nv

src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py

@@ -291,15 +291,15 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MambaModelProvid
         if kv_channels is not None:
             provider.kv_channels = kv_channels
 
-            provider.moe_aux_loss_coeff = 0.0001
-            provider.moe_router_score_function = "sigmoid"
-            provider.moe_router_enable_expert_bias = True
-            provider.moe_router_load_balancing_type = "seq_aux_loss"
-            provider.moe_router_dtype = "fp32"
-            provider.moe_grouped_gemm = True
-            provider.moe_token_dispatcher_type = "alltoall"
-            provider.moe_permute_fusion = True
-            provider.moe_shared_expert_overlap = True
+        provider.moe_aux_loss_coeff = 0.0001
+        provider.moe_router_score_function = "sigmoid"
+        provider.moe_router_enable_expert_bias = True
+        provider.moe_router_load_balancing_type = "seq_aux_loss"
+        provider.moe_router_dtype = "fp32"
+        provider.moe_grouped_gemm = True
+        provider.moe_token_dispatcher_type = "alltoall"
+        provider.moe_permute_fusion = True
+        provider.moe_shared_expert_overlap = True
 
         if hasattr(hf_config, "moe_latent_size"):
             provider.moe_latent_size = hf_config.moe_latent_size

src/megatron/bridge/recipes/nemotronh/nemotron_3_super.py

@@ -66,7 +66,6 @@ def nemotron_3_super_pretrain_config() -> ConfigContainer:
     cfg.model.moe_token_dispatcher_type = "alltoall"
     cfg.model.moe_shared_expert_overlap = False
     cfg.model.moe_flex_dispatcher_backend = "hybridep"
-    cfg.model.moe_router_force_load_balancing = True
 
     # Training Configuration
     cfg.train.train_iters = 39735
@@ -190,7 +189,6 @@ def nemotron_3_super_finetune_config(
     cfg.model.moe_token_dispatcher_type = "alltoall"
     cfg.model.moe_shared_expert_overlap = False
     cfg.model.moe_flex_dispatcher_backend = "hybridep"
-    cfg.model.moe_router_force_load_balancing = True
 
     # CUDA Graph (TE impl + partial scopes: ~40% throughput gain over disabled)
     cfg.model.cuda_graph_impl = "transformer_engine"
@@ -290,7 +288,6 @@ def nemotron_3_super_debug_pretrain_config() -> ConfigContainer:
     cfg.model.moe_token_dispatcher_type = "alltoall"
     cfg.model.moe_shared_expert_overlap = False
     cfg.model.moe_flex_dispatcher_backend = "hybridep"
-    cfg.model.moe_router_force_load_balancing = True
 
     # CUDA Graph (TE impl + partial scopes: ~40% throughput gain over disabled)
     cfg.model.cuda_graph_impl = "transformer_engine"