diff --git a/docs/guides/vlm/nemotron-omni.md b/docs/guides/vlm/nemotron-omni.md index 2ca682c306..eccd3c83dd 100644 --- a/docs/guides/vlm/nemotron-omni.md +++ b/docs/guides/vlm/nemotron-omni.md @@ -140,7 +140,8 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - enable_deepep: false + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/audio_finetune/qwen3_omni_asr/ami_sft.yaml b/examples/audio_finetune/qwen3_omni_asr/ami_sft.yaml index 91225e90a6..09de0ca3b4 100644 --- a/examples/audio_finetune/qwen3_omni_asr/ami_sft.yaml +++ b/examples/audio_finetune/qwen3_omni_asr/ami_sft.yaml @@ -66,8 +66,8 @@ model: linear: te rms_norm: torch_fp32 rope_fusion: false - experts: te - enable_deepep: true + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/audio_finetune/qwen3_omni_asr/multi_en_sft.yaml b/examples/audio_finetune/qwen3_omni_asr/multi_en_sft.yaml index bad1787bd2..c62f9216a1 100644 --- a/examples/audio_finetune/qwen3_omni_asr/multi_en_sft.yaml +++ b/examples/audio_finetune/qwen3_omni_asr/multi_en_sft.yaml @@ -55,8 +55,8 @@ model: linear: te rms_norm: torch_fp32 rope_fusion: false - experts: te - enable_deepep: true + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/llm_benchmark/glm/glm_4.7_te_deepep.yaml b/examples/llm_benchmark/glm/glm_4.7_te_deepep.yaml index f5ac8b1912..a3bf1fe426 100644 --- a/examples/llm_benchmark/glm/glm_4.7_te_deepep.yaml +++ b/examples/llm_benchmark/glm/glm_4.7_te_deepep.yaml @@ -79,7 +79,7 @@ model: rms_norm: te rope_fusion: true experts: gmm - dispatcher: deepep + dispatcher: hybridep fake_balanced_gate: true enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_benchmark/qwen/qwen3_next_te_deepep.yaml b/examples/llm_benchmark/qwen/qwen3_next_te_deepep.yaml index 2decc69ae4..f680a0ee64 100644 --- a/examples/llm_benchmark/qwen/qwen3_next_te_deepep.yaml +++ b/examples/llm_benchmark/qwen/qwen3_next_te_deepep.yaml @@ -78,7 +78,7 @@ model: rms_norm: te rope_fusion: true experts: gmm - dispatcher: deepep + dispatcher: hybridep fake_balanced_gate: true enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/deepseek_v32/deepseek_v32_hellaswag_pp.yaml b/examples/llm_finetune/deepseek_v32/deepseek_v32_hellaswag_pp.yaml index 1368c7a7ab..124a8a831c 100644 --- a/examples/llm_finetune/deepseek_v32/deepseek_v32_hellaswag_pp.yaml +++ b/examples/llm_finetune/deepseek_v32/deepseek_v32_hellaswag_pp.yaml @@ -73,7 +73,8 @@ model: linear: te rms_norm: torch_fp32 rope_fusion: false - enable_deepep: true + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag.yaml b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag.yaml index c52655bdbb..b8af2bba7e 100644 --- a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag.yaml +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag.yaml @@ -76,7 +76,7 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - dispatcher: deepep + dispatcher: hybridep experts: torch_mm enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml index 7ebb18a80a..c5c1ce9758 100644 --- a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml @@ -63,7 +63,7 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - dispatcher: deepep + dispatcher: hybridep experts: torch_mm enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_packed_sequence_hellaswag.yaml b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_packed_sequence_hellaswag.yaml index 7454d361c4..ad518eb6c2 100644 --- a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_packed_sequence_hellaswag.yaml +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_packed_sequence_hellaswag.yaml @@ -71,7 +71,7 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - dispatcher: deepep + dispatcher: hybridep experts: torch_mm enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/deepseek_v4/deepseek_v4_pro_hellaswag_all_tilelang_pp8_ep64_20steps.yaml b/examples/llm_finetune/deepseek_v4/deepseek_v4_pro_hellaswag_all_tilelang_pp8_ep64_20steps.yaml index b6a5611fbd..f9fc1501b1 100644 --- a/examples/llm_finetune/deepseek_v4/deepseek_v4_pro_hellaswag_all_tilelang_pp8_ep64_20steps.yaml +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_pro_hellaswag_all_tilelang_pp8_ep64_20steps.yaml @@ -73,7 +73,7 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - dispatcher: deepep + dispatcher: hybridep experts: torch_mm enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/deepseek_v4/deepseek_v4_pro_packed_sequence_hellaswag_all_tilelang_pp8_ep64_1024_20steps.yaml b/examples/llm_finetune/deepseek_v4/deepseek_v4_pro_packed_sequence_hellaswag_all_tilelang_pp8_ep64_1024_20steps.yaml index 5f69c5f734..526690bc18 100644 --- a/examples/llm_finetune/deepseek_v4/deepseek_v4_pro_packed_sequence_hellaswag_all_tilelang_pp8_ep64_1024_20steps.yaml +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_pro_packed_sequence_hellaswag_all_tilelang_pp8_ep64_1024_20steps.yaml @@ -73,7 +73,7 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - dispatcher: deepep + dispatcher: hybridep experts: torch_mm enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/glm/glm_4.7_te_deepep.yaml b/examples/llm_finetune/glm/glm_4.7_te_deepep.yaml index fc0e89881d..bcd5a113d1 100644 --- a/examples/llm_finetune/glm/glm_4.7_te_deepep.yaml +++ b/examples/llm_finetune/glm/glm_4.7_te_deepep.yaml @@ -37,6 +37,9 @@ rng: model: _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained pretrained_model_name_or_path: zai-org/GLM-4.7 + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + dispatcher: hybridep checkpoint: enabled: false diff --git a/examples/llm_finetune/glm/glm_5.1_hellaswag_pp.yaml b/examples/llm_finetune/glm/glm_5.1_hellaswag_pp.yaml index 1a8c61f857..62c24899a2 100644 --- a/examples/llm_finetune/glm/glm_5.1_hellaswag_pp.yaml +++ b/examples/llm_finetune/glm/glm_5.1_hellaswag_pp.yaml @@ -49,7 +49,8 @@ model: linear: te rms_norm: torch_fp32 rope_fusion: false - enable_deepep: true + experts: gmm + dispatcher: hybridep fake_balanced_gate: false enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/glm/glm_5.1_lora.yaml b/examples/llm_finetune/glm/glm_5.1_lora.yaml index a20bce7138..065a938c83 100644 --- a/examples/llm_finetune/glm/glm_5.1_lora.yaml +++ b/examples/llm_finetune/glm/glm_5.1_lora.yaml @@ -28,7 +28,7 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - dispatcher: deepep + dispatcher: hybridep experts: gmm fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/llm_finetune/glm/glm_5_hellaswag_pp.yaml b/examples/llm_finetune/glm/glm_5_hellaswag_pp.yaml index 27ab4e9c3e..5f03d28123 100644 --- a/examples/llm_finetune/glm/glm_5_hellaswag_pp.yaml +++ b/examples/llm_finetune/glm/glm_5_hellaswag_pp.yaml @@ -62,7 +62,8 @@ model: linear: te rms_norm: torch_fp32 rope_fusion: false - enable_deepep: true + experts: gmm + dispatcher: hybridep fake_balanced_gate: false enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft.yaml b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft.yaml index c0fb10ab78..5bc318a8a7 100644 --- a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft.yaml +++ b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft.yaml @@ -15,7 +15,8 @@ model: attn_implementation: sdpa backend: _target_: nemo_automodel.components.models.common.utils.BackendConfig - enable_deepep: false + experts: gmm + dispatcher: deepep distributed: _target_: nemo_automodel.components.distributed.fsdp2.FSDP2Manager dp_size: 8 diff --git a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft_chat.yaml b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft_chat.yaml index 0515ba5379..010ea363ce 100644 --- a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft_chat.yaml +++ b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft_chat.yaml @@ -15,7 +15,8 @@ model: attn_implementation: sdpa backend: _target_: nemo_automodel.components.models.common.utils.BackendConfig - enable_deepep: false + experts: gmm + dispatcher: deepep distributed: _target_: nemo_automodel.components.distributed.fsdp2.FSDP2Manager dp_size: 8 diff --git a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft.yaml b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft.yaml index 15ff30d4ab..2820598193 100644 --- a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft.yaml +++ b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft.yaml @@ -15,7 +15,7 @@ model: attn_implementation: sdpa backend: _target_: nemo_automodel.components.models.common.utils.BackendConfig - enable_deepep: false + dispatcher: torch distributed: _target_: nemo_automodel.components.distributed.fsdp2.FSDP2Manager dp_size: 1 diff --git a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft_packing.yaml b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft_packing.yaml index 6795d9248f..a3f3d54354 100644 --- a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft_packing.yaml +++ b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft_packing.yaml @@ -15,7 +15,7 @@ model: attn_implementation: sdpa backend: _target_: nemo_automodel.components.models.common.utils.BackendConfig - enable_deepep: false + dispatcher: torch distributed: _target_: nemo_automodel.components.distributed.fsdp2.FSDP2Manager dp_size: 1 diff --git a/examples/llm_finetune/hy_v3/hy3_preview_deepep_lora.yaml b/examples/llm_finetune/hy_v3/hy3_preview_deepep_lora.yaml index ad25bf6623..21607b4ead 100644 --- a/examples/llm_finetune/hy_v3/hy3_preview_deepep_lora.yaml +++ b/examples/llm_finetune/hy_v3/hy3_preview_deepep_lora.yaml @@ -46,7 +46,7 @@ model: linear: torch rms_norm: torch_fp32 experts: torch_mm - dispatcher: deepep + dispatcher: hybridep fake_balanced_gate: false gate_precision: float32 enable_hf_state_dict_adapter: true diff --git a/examples/llm_finetune/ling/ling_flash_2_0_sft.yaml b/examples/llm_finetune/ling/ling_flash_2_0_sft.yaml index 326336d580..3e5928bc13 100644 --- a/examples/llm_finetune/ling/ling_flash_2_0_sft.yaml +++ b/examples/llm_finetune/ling/ling_flash_2_0_sft.yaml @@ -42,7 +42,7 @@ model: linear: torch rms_norm: torch_fp32 experts: torch_mm - dispatcher: deepep + dispatcher: hybridep fake_balanced_gate: false enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/mimo_v2_flash/mimo_v2_flash_hellaswag.yaml b/examples/llm_finetune/mimo_v2_flash/mimo_v2_flash_hellaswag.yaml index 4e5d3a6941..9d4c62c7df 100644 --- a/examples/llm_finetune/mimo_v2_flash/mimo_v2_flash_hellaswag.yaml +++ b/examples/llm_finetune/mimo_v2_flash/mimo_v2_flash_hellaswag.yaml @@ -68,7 +68,7 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - dispatcher: deepep + dispatcher: hybridep experts: torch_mm gate_precision: float32 enable_hf_state_dict_adapter: true diff --git a/examples/llm_finetune/minimax_m2/minimax_m2.1_hellaswag_pp.yaml b/examples/llm_finetune/minimax_m2/minimax_m2.1_hellaswag_pp.yaml index 623c79e333..556aa2461e 100644 --- a/examples/llm_finetune/minimax_m2/minimax_m2.1_hellaswag_pp.yaml +++ b/examples/llm_finetune/minimax_m2/minimax_m2.1_hellaswag_pp.yaml @@ -63,7 +63,7 @@ model: linear: te rms_norm: torch_fp32 rope_fusion: true - dispatcher: deepep + dispatcher: hybridep experts: gmm fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/llm_finetune/minimax_m2/minimax_m2.5_hellaswag_pp.yaml b/examples/llm_finetune/minimax_m2/minimax_m2.5_hellaswag_pp.yaml index cb295b885c..7f4c8e604e 100644 --- a/examples/llm_finetune/minimax_m2/minimax_m2.5_hellaswag_pp.yaml +++ b/examples/llm_finetune/minimax_m2/minimax_m2.5_hellaswag_pp.yaml @@ -66,7 +66,7 @@ model: rms_norm: torch_fp32 rope_fusion: true experts: torch_mm - dispatcher: deepep + dispatcher: hybridep fake_balanced_gate: false enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_lora.yaml b/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_lora.yaml index 621b7e1260..c0b0115b6d 100644 --- a/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_lora.yaml +++ b/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_lora.yaml @@ -53,7 +53,7 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: true - dispatcher: deepep + dispatcher: hybridep experts: gmm fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_pp.yaml b/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_pp.yaml index 6a2d8eea54..f675b03c22 100644 --- a/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_pp.yaml +++ b/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_pp.yaml @@ -66,7 +66,7 @@ model: rms_norm: torch_fp32 rope_fusion: false experts: torch_mm - dispatcher: deepep + dispatcher: hybridep fake_balanced_gate: false enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft.yaml b/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft.yaml index 78aa42b59a..fcd6bc1a20 100644 --- a/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft.yaml +++ b/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft.yaml @@ -15,7 +15,8 @@ model: attn_implementation: sdpa backend: _target_: nemo_automodel.components.models.common.utils.BackendConfig - enable_deepep: false + experts: gmm + dispatcher: deepep distributed: _target_: nemo_automodel.components.distributed.fsdp2.FSDP2Manager dp_size: 8 diff --git a/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft_chat.yaml b/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft_chat.yaml index 0ef677a9b9..7392246569 100644 --- a/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft_chat.yaml +++ b/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft_chat.yaml @@ -15,7 +15,8 @@ model: attn_implementation: sdpa backend: _target_: nemo_automodel.components.models.common.utils.BackendConfig - enable_deepep: false + experts: gmm + dispatcher: deepep distributed: _target_: nemo_automodel.components.distributed.fsdp2.FSDP2Manager dp_size: 8 diff --git a/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft.yaml b/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft.yaml index 78cb713402..035b449953 100644 --- a/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft.yaml +++ b/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft.yaml @@ -15,7 +15,8 @@ model: attn_implementation: sdpa backend: _target_: nemo_automodel.components.models.common.utils.BackendConfig - enable_deepep: false + experts: gmm + dispatcher: deepep distributed: _target_: nemo_automodel.components.distributed.fsdp2.FSDP2Manager dp_size: null diff --git a/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft_packing.yaml b/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft_packing.yaml index ce894a09e7..7ab97b3c73 100644 --- a/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft_packing.yaml +++ b/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft_packing.yaml @@ -15,7 +15,8 @@ model: attn_implementation: sdpa backend: _target_: nemo_automodel.components.models.common.utils.BackendConfig - enable_deepep: false + experts: gmm + dispatcher: deepep distributed: _target_: nemo_automodel.components.distributed.fsdp2.FSDP2Manager dp_size: null diff --git a/examples/llm_finetune/nemotron/nemotron_ultra_v3_hellaswag_peft.yaml b/examples/llm_finetune/nemotron/nemotron_ultra_v3_hellaswag_peft.yaml index e676c53dcc..8f8f0a16f6 100644 --- a/examples/llm_finetune/nemotron/nemotron_ultra_v3_hellaswag_peft.yaml +++ b/examples/llm_finetune/nemotron/nemotron_ultra_v3_hellaswag_peft.yaml @@ -63,7 +63,7 @@ model: linear: te rms_norm: te experts: gmm - dispatcher: deepep + dispatcher: hybridep fake_balanced_gate: false enable_fsdp_optimizations: false diff --git a/examples/llm_finetune/nemotron/nemotron_ultra_v3_hellaswag_peft_gb200.yaml b/examples/llm_finetune/nemotron/nemotron_ultra_v3_hellaswag_peft_gb200.yaml index 499dace967..0821e87ecf 100644 --- a/examples/llm_finetune/nemotron/nemotron_ultra_v3_hellaswag_peft_gb200.yaml +++ b/examples/llm_finetune/nemotron/nemotron_ultra_v3_hellaswag_peft_gb200.yaml @@ -133,3 +133,8 @@ lr_scheduler: lr_decay_style: cosine lr_warmup_steps: 10 min_lr: 1.0e-5 + +ci: + # pp_size(1) * ep_size(16) = 16 GPUs => 4 nodes (4 GB200/node). + recipe_owner: adil-a + nodes: 4 diff --git a/examples/llm_finetune/nemotron/nemotron_ultra_v3_squad.yaml b/examples/llm_finetune/nemotron/nemotron_ultra_v3_squad.yaml index a0fa50eaae..c00b0a9d99 100644 --- a/examples/llm_finetune/nemotron/nemotron_ultra_v3_squad.yaml +++ b/examples/llm_finetune/nemotron/nemotron_ultra_v3_squad.yaml @@ -121,6 +121,11 @@ lr_scheduler: lr_warmup_steps: 10 min_lr: 1.0e-6 +ci: + # pp_size(1) * ep_size(64) = 64 GPUs => 8 nodes (8 H100/node). + recipe_owner: adil-a + nodes: 8 + # wandb: # project: my-project # entity: my-entity diff --git a/examples/llm_finetune/qwen/qwen3_next_te_deepep.yaml b/examples/llm_finetune/qwen/qwen3_next_te_deepep.yaml index a53aafc78d..42f64ad56b 100644 --- a/examples/llm_finetune/qwen/qwen3_next_te_deepep.yaml +++ b/examples/llm_finetune/qwen/qwen3_next_te_deepep.yaml @@ -38,6 +38,9 @@ rng: model: _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained pretrained_model_name_or_path: Qwen/Qwen3-Next-80B-A3B-Instruct + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + dispatcher: hybridep checkpoint: enabled: false diff --git a/examples/llm_finetune/stepfun/step_3.5_flash_hellaswag_pp.yaml b/examples/llm_finetune/stepfun/step_3.5_flash_hellaswag_pp.yaml index 7192942154..de949b49d7 100644 --- a/examples/llm_finetune/stepfun/step_3.5_flash_hellaswag_pp.yaml +++ b/examples/llm_finetune/stepfun/step_3.5_flash_hellaswag_pp.yaml @@ -65,7 +65,8 @@ model: linear: te rms_norm: torch_fp32 rope_fusion: true - enable_deepep: true + experts: gmm + dispatcher: hybridep fake_balanced_gate: false enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/llm_pretrain/deepseekv3_pretrain.yaml b/examples/llm_pretrain/deepseekv3_pretrain.yaml index 6dae590ff5..bff57a442b 100644 --- a/examples/llm_pretrain/deepseekv3_pretrain.yaml +++ b/examples/llm_pretrain/deepseekv3_pretrain.yaml @@ -152,3 +152,8 @@ lr_scheduler: lr_decay_style: cosine lr_warmup_steps: 500 min_lr: 0.0 + +ci: + # pp_size(8) * ep_size(16) = 128 GPUs => 16 nodes (8 H100/node). + recipe_owner: hemildesai + nodes: 16 diff --git a/examples/vlm_finetune/kimi/kimi25vl_medpix.yaml b/examples/vlm_finetune/kimi/kimi25vl_medpix.yaml index 3c9257b7c1..f3f63f5cbe 100644 --- a/examples/vlm_finetune/kimi/kimi25vl_medpix.yaml +++ b/examples/vlm_finetune/kimi/kimi25vl_medpix.yaml @@ -62,7 +62,8 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - enable_deepep: true + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/vlm_finetune/kimi/kimi2vl_cordv2.yaml b/examples/vlm_finetune/kimi/kimi2vl_cordv2.yaml index 5fc11c377a..1abcdef428 100644 --- a/examples/vlm_finetune/kimi/kimi2vl_cordv2.yaml +++ b/examples/vlm_finetune/kimi/kimi2vl_cordv2.yaml @@ -61,7 +61,8 @@ model: linear: te rms_norm: torch_fp32 rope_fusion: true - enable_deepep: true + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/vlm_finetune/qwen3/qwen3_omni_moe_30b_te_deepep.yaml b/examples/vlm_finetune/qwen3/qwen3_omni_moe_30b_te_deepep.yaml index 8ad20de42f..11486c4f9b 100644 --- a/examples/vlm_finetune/qwen3/qwen3_omni_moe_30b_te_deepep.yaml +++ b/examples/vlm_finetune/qwen3/qwen3_omni_moe_30b_te_deepep.yaml @@ -45,8 +45,8 @@ model: linear: te rms_norm: torch_fp32 rope_fusion: false - experts: te - enable_deepep: true + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/vlm_finetune/qwen3/qwen3_vl_moe_235b.yaml b/examples/vlm_finetune/qwen3/qwen3_vl_moe_235b.yaml index fbaa242b24..ffb9602e56 100644 --- a/examples/vlm_finetune/qwen3/qwen3_vl_moe_235b.yaml +++ b/examples/vlm_finetune/qwen3/qwen3_vl_moe_235b.yaml @@ -46,7 +46,8 @@ model: linear: te rms_norm: torch_fp32 rope_fusion: false - enable_deepep: true + experts: gmm + dispatcher: hybridep fake_balanced_gate: false enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_neat_packing.yaml b/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_neat_packing.yaml index aabc91bd17..3e91749092 100644 --- a/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_neat_packing.yaml +++ b/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_neat_packing.yaml @@ -44,8 +44,8 @@ model: linear: te rms_norm: te rope_fusion: false - experts: te - enable_deepep: true + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_te_deepep.yaml b/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_te_deepep.yaml index ee2f7bff9e..2a6d1db089 100644 --- a/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_te_deepep.yaml +++ b/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_te_deepep.yaml @@ -47,8 +47,8 @@ model: linear: te rms_norm: torch_fp32 rope_fusion: false - experts: te - enable_deepep: true + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b.yaml b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b.yaml index bc0006a2ac..59c43f154f 100644 --- a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b.yaml +++ b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b.yaml @@ -44,7 +44,8 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - enable_deepep: true + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml index 7c9c326f77..cc19c1e258 100644 --- a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml +++ b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml @@ -49,7 +49,8 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - enable_deepep: true + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_moe_medpix.yaml b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_moe_medpix.yaml index 275bd0ec5d..cc7dd9fe26 100644 --- a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_moe_medpix.yaml +++ b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_moe_medpix.yaml @@ -45,7 +45,8 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - enable_deepep: false + experts: gmm + dispatcher: hybridep fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b.yaml b/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b.yaml index dc5a3d9066..9455cda80a 100644 --- a/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b.yaml +++ b/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b.yaml @@ -48,7 +48,8 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - enable_deepep: true + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml b/examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml index eb913a679a..020dd7613c 100644 --- a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml +++ b/examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml @@ -45,7 +45,7 @@ model: rms_norm: torch_fp32 rope_fusion: false experts: gmm - dispatcher: deepep + dispatcher: hybridep fake_balanced_gate: false enable_hf_state_dict_adapter: true enable_fsdp_optimizations: true diff --git a/fern/versions/nightly/pages/guides/vlm/nemotron-omni.mdx b/fern/versions/nightly/pages/guides/vlm/nemotron-omni.mdx index 30ec1e7644..704f847b39 100644 --- a/fern/versions/nightly/pages/guides/vlm/nemotron-omni.mdx +++ b/fern/versions/nightly/pages/guides/vlm/nemotron-omni.mdx @@ -144,7 +144,8 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - enable_deepep: false + experts: gmm + dispatcher: deepep fake_balanced_gate: false enable_hf_state_dict_adapter: true diff --git a/nemo_automodel/components/models/common/utils.py b/nemo_automodel/components/models/common/utils.py index 72d51a4fb4..e3ba879056 100644 --- a/nemo_automodel/components/models/common/utils.py +++ b/nemo_automodel/components/models/common/utils.py @@ -14,7 +14,6 @@ import importlib.util import logging -import warnings from contextlib import contextmanager, nullcontext from dataclasses import dataclass from typing import Any, Literal @@ -176,7 +175,8 @@ class BackendConfig: manager instance across MoE layers. dispatcher_async_dispatch: Whether DeepEP/UCCL-EP dispatch should return asynchronously and allocate dispatched tensors on the communication stream. - enable_deepep: Deprecated. Use dispatcher="deepep" and experts="gmm" instead. + enable_deepep: Removed and ignored. Logs a warning if set; configure "dispatcher" + and "experts" explicitly instead. fake_balanced_gate: If True, replace the learned Gate with FakeBalancedGate that assigns tokens to experts without learned routing weights. fake_gate_noise: Noise level [0, 1] for FakeBalancedGate. When > 0, uses @@ -210,7 +210,7 @@ class BackendConfig: dispatcher_num_sms: int = 20 dispatcher_share_token_dispatcher: bool = True dispatcher_async_dispatch: bool = False - enable_deepep: bool | None = None # Deprecated: use dispatcher="deepep" instead + enable_deepep: bool | None = None # Removed: ignored with a warning; set dispatcher/experts explicitly fake_balanced_gate: bool = False # Approximate max/mean load ratios (64 experts, top-8, 4096 tokens): # 0.0→1.00x, 0.1→~1.2x, 0.3→~1.6x, 0.5→~2.0x, 1.0→~2.8x. @@ -235,21 +235,17 @@ def __post_init__(self): if isinstance(self.gate_precision, str): self.gate_precision = dtype_from_str(self.gate_precision, default=None) - # Handle deprecated enable_deepep parameter + # enable_deepep was removed. It is no longer honored; warn (once, on rank 0) if a stale + # config still sets it so the user migrates to explicit dispatcher/experts. The field is + # retained only so loading an old config does not crash this kw_only dataclass. if self.enable_deepep is not None: - warnings.warn( - "enable_deepep is deprecated and will be removed in a future release. " - "Use experts='gmm' and dispatcher='deepep' instead of enable_deepep=True, " - "or dispatcher='torch' instead of enable_deepep=False.", - DeprecationWarning, - stacklevel=2, - ) - if self.enable_deepep: - self.experts = "gmm" - self.dispatcher = "deepep" - else: - self.dispatcher = "torch" - # Clear the deprecated field after conversion + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + logger.warning( + "enable_deepep is no longer supported and is ignored. " + "Set 'dispatcher' (deepep/hybridep/torch) and 'experts' explicitly instead. " + "Previously enable_deepep=True was equivalent to experts=gmm + dispatcher=deepep, " + "and enable_deepep=False to dispatcher=torch." + ) self.enable_deepep = None # Backward compatibility diff --git a/tests/functional_tests/hf_transformer_finetune/nemotron_nano_v3_4layer_custom.yaml b/tests/functional_tests/hf_transformer_finetune/nemotron_nano_v3_4layer_custom.yaml index 884e33f264..d4e62ce36b 100644 --- a/tests/functional_tests/hf_transformer_finetune/nemotron_nano_v3_4layer_custom.yaml +++ b/tests/functional_tests/hf_transformer_finetune/nemotron_nano_v3_4layer_custom.yaml @@ -36,7 +36,7 @@ model: trust_remote_code: true backend: _target_: nemo_automodel.components.models.common.BackendConfig - enable_deepep: false + dispatcher: torch checkpoint: enabled: false diff --git a/tests/unit_tests/models/deepseek_v32/test_dsv32_state_dict_adapter.py b/tests/unit_tests/models/deepseek_v32/test_dsv32_state_dict_adapter.py index 6cfafcab61..b19f404883 100644 --- a/tests/unit_tests/models/deepseek_v32/test_dsv32_state_dict_adapter.py +++ b/tests/unit_tests/models/deepseek_v32/test_dsv32_state_dict_adapter.py @@ -12,27 +12,27 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib.util import sys import types -import importlib.util -import pytest +from unittest.mock import Mock, patch + import torch -from unittest.mock import Mock, patch, MagicMock # Mock fast_hadamard_transform before importing deepseek_v32 modules try: import fast_hadamard_transform # noqa: F401 except ImportError: - if 'fast_hadamard_transform' not in sys.modules: - mock_hadamard = types.ModuleType('fast_hadamard_transform') - mock_hadamard.__spec__ = importlib.util.spec_from_loader('fast_hadamard_transform', loader=None) + if "fast_hadamard_transform" not in sys.modules: + mock_hadamard = types.ModuleType("fast_hadamard_transform") + mock_hadamard.__spec__ = importlib.util.spec_from_loader("fast_hadamard_transform", loader=None) mock_hadamard.hadamard_transform = lambda x, scale: x - sys.modules['fast_hadamard_transform'] = mock_hadamard + sys.modules["fast_hadamard_transform"] = mock_hadamard +from nemo_automodel.components.models.common import BackendConfig from nemo_automodel.components.models.deepseek_v32.config import DeepseekV32Config from nemo_automodel.components.models.deepseek_v32.state_dict_adapter import DeepSeekV32StateDictAdapter from nemo_automodel.components.moe.config import MoEConfig -from nemo_automodel.components.models.common import BackendConfig class TestDeepSeekV32StateDictAdapter: @@ -62,7 +62,7 @@ def create_mock_moe_config(self, **overrides): def create_mock_backend_config(self, **overrides): backend = Mock(spec=BackendConfig) - backend.enable_deepep = False + backend.dispatcher = "torch" for key, value in overrides.items(): setattr(backend, key, value) @@ -76,10 +76,7 @@ def test_initialization(self): backend = self.create_mock_backend_config() adapter = DeepSeekV32StateDictAdapter( - config=config, - moe_config=moe_config, - backend=backend, - dtype=torch.float16 + config=config, moe_config=moe_config, backend=backend, dtype=torch.float16 ) assert adapter.config == config @@ -138,7 +135,7 @@ def create_mock_moe_config(self, **overrides): def create_mock_backend_config(self, **overrides): backend = Mock(spec=BackendConfig) - backend.enable_deepep = False + backend.dispatcher = "torch" for key, value in overrides.items(): setattr(backend, key, value) @@ -156,7 +153,7 @@ def test_convert_tensor_quantization_normal_weight(self): tensor = torch.randn(256, 128) fqn = "model.layers.0.self_attn.q_a_proj.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn, tensor, quantization=True) assert len(result) == 2 @@ -176,7 +173,7 @@ def test_convert_tensor_quantization_skips_layernorm(self): tensor = torch.randn(256) fqn = "model.layers.0.input_layernorm.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn, tensor, quantization=True) assert len(result) == 1 @@ -195,7 +192,7 @@ def test_convert_tensor_quantization_skips_indexer_k_norm(self): tensor_weight = torch.randn(64) fqn_weight = "model.layers.0.self_attn.indexer.k_norm.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn_weight, tensor_weight, quantization=True) assert len(result) == 1 @@ -206,7 +203,7 @@ def test_convert_tensor_quantization_skips_indexer_k_norm(self): tensor_bias = torch.randn(64) fqn_bias = "model.layers.0.self_attn.indexer.k_norm.bias" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn_bias, tensor_bias, quantization=True) assert len(result) == 1 @@ -225,7 +222,7 @@ def test_convert_tensor_quantization_indexer_linear_weights(self): tensor = torch.randn(256, 128) fqn = "model.layers.0.self_attn.indexer.wq_b.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn, tensor, quantization=True) assert len(result) == 2 @@ -237,7 +234,7 @@ def test_convert_tensor_quantization_indexer_linear_weights(self): tensor_wk = torch.randn(256, 128) fqn_wk = "model.layers.0.self_attn.indexer.wk.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn_wk, tensor_wk, quantization=True) assert len(result) == 2 @@ -257,7 +254,7 @@ def test_convert_tensor_quantization_skips_indexer_weights_proj(self): tensor = torch.randn(256, 128) fqn = "model.layers.0.self_attn.indexer.weights_proj.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn, tensor, quantization=True) assert len(result) == 1 @@ -292,7 +289,7 @@ def create_mock_moe_config(self, **overrides): def create_mock_backend_config(self, **overrides): backend = Mock(spec=BackendConfig) - backend.enable_deepep = False + backend.dispatcher = "torch" for key, value in overrides.items(): setattr(backend, key, value) @@ -394,7 +391,7 @@ def create_mock_moe_config(self, **overrides): def create_mock_backend_config(self, **overrides): backend = Mock(spec=BackendConfig) - backend.enable_deepep = False + backend.dispatcher = "torch" for key, value in overrides.items(): setattr(backend, key, value) @@ -412,7 +409,7 @@ def test_convert_tensor_with_exclude_regex(self): tensor = torch.randn(256, 128) fqn = "model.layers.0.self_attn.q_a_proj.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): # With regex that matches the key result = adapter.convert_single_tensor_to_hf( fqn, tensor, quantization=False, exclude_key_regex=r".*q_a_proj.*" @@ -432,7 +429,7 @@ def test_convert_tensor_with_exclude_regex_no_match(self): tensor = torch.randn(256, 128) fqn = "model.layers.0.self_attn.q_a_proj.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): # With regex that doesn't match the key result = adapter.convert_single_tensor_to_hf( fqn, tensor, quantization=False, exclude_key_regex=r".*kv_proj.*" @@ -453,7 +450,7 @@ def test_convert_tensor_without_quantization(self): tensor = torch.randn(256, 128) fqn = "model.layers.0.self_attn.q_a_proj.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn, tensor, quantization=False) # Without quantization, should just return the tensor @@ -491,7 +488,7 @@ def create_mock_moe_config(self, **overrides): def create_mock_backend_config(self, **overrides): backend = Mock(spec=BackendConfig) - backend.enable_deepep = False + backend.dispatcher = "torch" for key, value in overrides.items(): setattr(backend, key, value) @@ -555,7 +552,7 @@ def create_mock_moe_config(self, **overrides): def create_mock_backend_config(self, **overrides): backend = Mock(spec=BackendConfig) - backend.enable_deepep = False + backend.dispatcher = "torch" for key, value in overrides.items(): setattr(backend, key, value) @@ -565,8 +562,8 @@ def create_mock_backend_config(self, **overrides): def test_scale_shape_calculation(self): """Test that scale shape is calculated correctly.""" from nemo_automodel.components.models.deepseek_v3.state_dict_adapter import ( - calculate_scale_shape, BLOCK_SIZE, + calculate_scale_shape, ) # Test with tensor that divides evenly by block size @@ -588,7 +585,7 @@ def test_quantized_weights_have_correct_scale_shape(self): tensor = torch.randn(256, 128) fqn = "model.layers.0.self_attn.q_a_proj.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn, tensor, quantization=True) # Should have weight and scale_inv diff --git a/tests/unit_tests/models/minimax_m2/test_minimax_m2_layers.py b/tests/unit_tests/models/minimax_m2/test_minimax_m2_layers.py index 90212fe05e..24fc3bb18a 100644 --- a/tests/unit_tests/models/minimax_m2/test_minimax_m2_layers.py +++ b/tests/unit_tests/models/minimax_m2/test_minimax_m2_layers.py @@ -31,7 +31,7 @@ def backend(): attn="sdpa", rms_norm="torch", rope_fusion=False, - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) diff --git a/tests/unit_tests/models/minimax_m2/test_minimax_m2_model.py b/tests/unit_tests/models/minimax_m2/test_minimax_m2_model.py index 35f1b3785f..ccfe19663c 100644 --- a/tests/unit_tests/models/minimax_m2/test_minimax_m2_model.py +++ b/tests/unit_tests/models/minimax_m2/test_minimax_m2_model.py @@ -49,7 +49,7 @@ def backend(): attn="sdpa", rms_norm="torch", rope_fusion=False, - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) diff --git a/tests/unit_tests/models/minimax_m2/test_minimax_m2_state_dict_adapter.py b/tests/unit_tests/models/minimax_m2/test_minimax_m2_state_dict_adapter.py index 13597b0bfc..80f7f014b8 100644 --- a/tests/unit_tests/models/minimax_m2/test_minimax_m2_state_dict_adapter.py +++ b/tests/unit_tests/models/minimax_m2/test_minimax_m2_state_dict_adapter.py @@ -54,7 +54,7 @@ def backend(): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=True, ) diff --git a/tests/unit_tests/models/minimax_m2/test_minimax_m2_state_dict_adapter_dtensor.py b/tests/unit_tests/models/minimax_m2/test_minimax_m2_state_dict_adapter_dtensor.py index 7c12685891..f3cdd0bd89 100644 --- a/tests/unit_tests/models/minimax_m2/test_minimax_m2_state_dict_adapter_dtensor.py +++ b/tests/unit_tests/models/minimax_m2/test_minimax_m2_state_dict_adapter_dtensor.py @@ -46,7 +46,7 @@ def adapter(): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=True, ) diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_layers.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_layers.py index e73aabc5b8..8d316139df 100644 --- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_layers.py +++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_layers.py @@ -216,7 +216,7 @@ def backend(self): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) @@ -517,7 +517,7 @@ def backend(self): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_model.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_model.py index ed31d1e459..85c45aa31d 100644 --- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_model.py +++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_model.py @@ -97,7 +97,7 @@ def backend(self): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) @@ -328,7 +328,7 @@ def backend(self): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) @@ -762,7 +762,7 @@ def backend(self): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) @@ -960,7 +960,7 @@ def backend(self): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) @@ -1118,7 +1118,7 @@ def backend(self): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=True, # Use fake balanced gate for deterministic testing enable_hf_state_dict_adapter=False, ) diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp.py index 4c4ffecdf5..2f6644e02d 100644 --- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp.py +++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp.py @@ -223,7 +223,7 @@ def backend(): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=True, enable_hf_state_dict_adapter=False, ) diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp_parity.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp_parity.py index 14c0151d26..13c1827948 100644 --- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp_parity.py +++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp_parity.py @@ -48,7 +48,7 @@ def backend(): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=True, enable_hf_state_dict_adapter=True, ) diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_pp_mtp.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_pp_mtp.py index 0fdb3fc61d..316cd413a5 100644 --- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_pp_mtp.py +++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_pp_mtp.py @@ -37,7 +37,7 @@ def backend(): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=True, enable_hf_state_dict_adapter=False, ) diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py index 3e7d34e10d..b38d28caf1 100644 --- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py +++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py @@ -74,7 +74,7 @@ def backend(self): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", ) def test_adapter_init(self, config, moe_config, backend): @@ -228,7 +228,7 @@ def test_from_hf_renames_backbone_to_model(self, config, moe_config, backend): "lm_head.weight": torch.randn(100, 256), } - result = adapter.from_hf(hf_state_dict) + adapter.from_hf(hf_state_dict) # Check that _from_hf_w_merged_experts was called with renamed state dict call_args = mock_merge.call_args[0][0] @@ -486,7 +486,10 @@ def test_from_hf_uses_mixer_experts_path(self, config, moe_config, backend): with patch.object(adapter, "_validate_expert_availability"): with patch("nemo_automodel.components.moe.state_dict_mixin.should_load_expert_for_rank", return_value=True): - with patch("nemo_automodel.components.moe.state_dict_mixin.create_dtensor_from_local", side_effect=lambda x, *args: x): + with patch( + "nemo_automodel.components.moe.state_dict_mixin.create_dtensor_from_local", + side_effect=lambda x, *args: x, + ): result = adapter._from_hf_w_merged_experts(hf_state_dict) # Should have created merged expert tensors with mixer.experts path diff --git a/tests/unit_tests/models/qwen3_5/test_qwen3_5_mtp.py b/tests/unit_tests/models/qwen3_5/test_qwen3_5_mtp.py index 9baaf4c377..a7a91d7ab0 100644 --- a/tests/unit_tests/models/qwen3_5/test_qwen3_5_mtp.py +++ b/tests/unit_tests/models/qwen3_5/test_qwen3_5_mtp.py @@ -74,7 +74,7 @@ def _backend(): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=True, ) diff --git a/tests/unit_tests/models/qwen3_5/test_qwen3_5_pp.py b/tests/unit_tests/models/qwen3_5/test_qwen3_5_pp.py index 7203041735..5c1c79315c 100644 --- a/tests/unit_tests/models/qwen3_5/test_qwen3_5_pp.py +++ b/tests/unit_tests/models/qwen3_5/test_qwen3_5_pp.py @@ -194,7 +194,7 @@ def _tiny_vlm_model(): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=True, ) diff --git a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_model.py b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_model.py index d95ad4bc3c..cda01163ea 100644 --- a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_model.py +++ b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_model.py @@ -129,7 +129,7 @@ def backend_config(): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) diff --git a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_mtp.py b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_mtp.py index 2d03e8ed8e..71287b5fdb 100644 --- a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_mtp.py +++ b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_mtp.py @@ -66,7 +66,7 @@ def _backend(): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=True, ) diff --git a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_state_dict_adapter.py b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_state_dict_adapter.py index 77fa1cd85b..e936ce3810 100644 --- a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_state_dict_adapter.py +++ b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_state_dict_adapter.py @@ -70,7 +70,7 @@ def backend_config(): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) diff --git a/tests/unit_tests/models/step3p5/test_step3p5_layers.py b/tests/unit_tests/models/step3p5/test_step3p5_layers.py index 8b4f56f076..906d8dc8d0 100644 --- a/tests/unit_tests/models/step3p5/test_step3p5_layers.py +++ b/tests/unit_tests/models/step3p5/test_step3p5_layers.py @@ -26,13 +26,13 @@ Step3p5RotaryEmbedding, ) - pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @dataclass class MockStep3p5Config: """Mock configuration for Step3p5 model.""" + vocab_size: int = 128 hidden_size: int = 64 intermediate_size: int = 128 @@ -74,7 +74,7 @@ def sdpa_backend(): attn="sdpa", rms_norm="torch", rope_fusion=False, - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) @@ -277,7 +277,10 @@ def test_forward_shape_preserved(self, config, sdpa_backend): fake_attn = torch.zeros(batch, config.num_attention_heads, seq, config.head_dim) attention.attn_func = MagicMock(return_value=fake_attn.to(torch.bfloat16)) - with patch("nemo_automodel.components.models.step3p5.layers.apply_rotary_emb_qk", side_effect=lambda q, k, *_, **__: (q, k)): + with patch( + "nemo_automodel.components.models.step3p5.layers.apply_rotary_emb_qk", + side_effect=lambda q, k, *_, **__: (q, k), + ): out = attention(x, freqs_cis=freqs_cis) assert out.shape == (batch, seq, config.hidden_size) diff --git a/tests/unit_tests/models/step3p5/test_step3p5_model.py b/tests/unit_tests/models/step3p5/test_step3p5_model.py index 0d5e61f1c5..fcf6a2b278 100644 --- a/tests/unit_tests/models/step3p5/test_step3p5_model.py +++ b/tests/unit_tests/models/step3p5/test_step3p5_model.py @@ -92,7 +92,7 @@ def sdpa_backend(): attn="sdpa", rms_norm="torch", rope_fusion=False, - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False, ) diff --git a/tests/unit_tests/models/step3p5/test_step3p5_state_dict_adapter.py b/tests/unit_tests/models/step3p5/test_step3p5_state_dict_adapter.py index fb0c814518..0a5b699b87 100644 --- a/tests/unit_tests/models/step3p5/test_step3p5_state_dict_adapter.py +++ b/tests/unit_tests/models/step3p5/test_step3p5_state_dict_adapter.py @@ -70,7 +70,7 @@ def backend(): linear="torch", attn="sdpa", rms_norm="torch", - enable_deepep=False, + dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=True, ) diff --git a/tests/unit_tests/moe/test_backend_config.py b/tests/unit_tests/moe/test_backend_config.py index 529370cb40..a617239020 100644 --- a/tests/unit_tests/moe/test_backend_config.py +++ b/tests/unit_tests/moe/test_backend_config.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import warnings +import logging import pytest import torch @@ -122,66 +122,43 @@ def test_fake_gate_noise_with_fake_balanced_gate(self): assert config.fake_gate_noise == 0.3 -class TestBackendConfigEnableDeepepDeprecation: - """Test backwards compatibility for deprecated enable_deepep parameter.""" - - def test_enable_deepep_true_sets_dispatcher_deepep_and_experts_gmm(self): - """Test that enable_deepep=True sets dispatcher='deepep' and experts='gmm' with deprecation warning.""" - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - config = BackendConfig(enable_deepep=True) - assert config.dispatcher == "deepep" - assert config.experts == "gmm" - assert config.enable_deepep is None # Should be cleared after conversion - assert len(w) == 1 - assert issubclass(w[0].category, DeprecationWarning) - assert "enable_deepep is deprecated" in str(w[0].message) - - def test_enable_deepep_false_sets_dispatcher_torch(self): - """Test that enable_deepep=False sets dispatcher='torch' with deprecation warning.""" - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - config = BackendConfig(enable_deepep=False) - assert config.dispatcher == "torch" - expected_experts = "torch_mm" if torch.cuda.is_available() else "torch" - assert config.experts == expected_experts # experts unchanged when enable_deepep=False - assert config.enable_deepep is None # Should be cleared after conversion - assert len(w) == 1 - assert issubclass(w[0].category, DeprecationWarning) - assert "enable_deepep is deprecated" in str(w[0].message) - - def test_enable_deepep_none_no_warning(self): - """Test that enable_deepep=None (default) does not trigger warning.""" - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") +class TestBackendConfigEnableDeepepRemoved: + """enable_deepep was removed: it is ignored (with a warning) and never alters dispatcher/experts.""" + + def test_enable_deepep_true_is_ignored_and_warns(self, caplog): + """enable_deepep=True is ignored; dispatcher/experts keep their explicit values and a warning is logged.""" + with caplog.at_level(logging.WARNING): + config = BackendConfig(dispatcher="hybridep", experts="gmm", enable_deepep=True) + assert config.dispatcher == "hybridep" # not overridden to "deepep" + assert config.experts == "gmm" + assert config.enable_deepep is None # cleared after the warning + assert "enable_deepep is no longer supported" in caplog.text + + def test_enable_deepep_false_is_ignored_and_warns(self, caplog): + """enable_deepep=False is ignored; the dispatcher is NOT forced to torch and a warning is logged.""" + with caplog.at_level(logging.WARNING): + config = BackendConfig(dispatcher="deepep", experts="gmm", enable_deepep=False) + assert config.dispatcher == "deepep" # not forced to "torch" + assert config.experts == "gmm" + assert config.enable_deepep is None + assert "enable_deepep is no longer supported" in caplog.text + + def test_enable_deepep_none_no_warning(self, caplog): + """enable_deepep=None (default) leaves the field as None and logs no warning.""" + with caplog.at_level(logging.WARNING): config = BackendConfig() - assert config.enable_deepep is None - # No deprecation warning should be raised - deprecation_warnings = [x for x in w if issubclass(x.category, DeprecationWarning)] - assert len(deprecation_warnings) == 0 - - def test_enable_deepep_overrides_dispatcher_and_experts(self): - """Test that enable_deepep takes precedence over dispatcher and experts when both provided.""" - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - # Even if dispatcher="torch" and experts="torch", enable_deepep=True should override them + assert config.enable_deepep is None + assert "enable_deepep" not in caplog.text + + def test_enable_deepep_does_not_override_explicit_dispatcher(self, caplog): + """A stale enable_deepep no longer wins over an explicit dispatcher/experts.""" + with caplog.at_level(logging.WARNING): config = BackendConfig(dispatcher="torch", experts="torch", enable_deepep=True) - assert config.dispatcher == "deepep" - assert config.experts == "gmm" - assert len(w) == 1 - assert issubclass(w[0].category, DeprecationWarning) - - def test_enable_deepep_false_overrides_dispatcher_deepep(self): - """Test that enable_deepep=False overrides dispatcher='deepep'.""" - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - config = BackendConfig(dispatcher="deepep", enable_deepep=False) - assert config.dispatcher == "torch" - assert len(w) == 1 - assert issubclass(w[0].category, DeprecationWarning) + assert config.dispatcher == "torch" + assert config.experts == "torch" def test_dispatcher_without_enable_deepep(self): - """Test that dispatcher works correctly without enable_deepep.""" + """dispatcher works correctly without enable_deepep (field stays None).""" config = BackendConfig(dispatcher="deepep") assert config.dispatcher == "deepep" assert config.enable_deepep is None @@ -190,17 +167,6 @@ def test_dispatcher_without_enable_deepep(self): assert config.dispatcher == "torch" assert config.enable_deepep is None - def test_deprecation_warning_message_content(self): - """Test that deprecation warning message contains helpful migration info.""" - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - BackendConfig(enable_deepep=True) - warning_message = str(w[0].message) - assert "experts='gmm'" in warning_message - assert "dispatcher='deepep'" in warning_message - assert "dispatcher='torch'" in warning_message - assert "will be removed in a future release" in warning_message - class TestBackendConfigHybridEP: """Test BackendConfig HybridEP dispatcher support."""