[main] feat(moe): Support gated delta net for Qwen3-Next (1/4) (NVIDIA#1989)

Signed-off-by: oliver könig <[email protected]>
Co-authored-by: oliver könig <[email protected]>

Author: 2 people

gpt_builders.py

@@ -8,6 +8,9 @@
     get_gpt_layer_with_inference_spec,
     get_gpt_mtp_block_spec,
 )
+from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
+    get_transformer_block_with_experimental_attention_variant_spec,
+)
 from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import (
     get_gpt_heterogeneous_layer_spec,
 )
@@ -42,7 +45,13 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_
         else:
             use_te = args.transformer_impl == "transformer_engine"
 
-            if args.num_experts:
+            if args.experimental_attention_variant is not None:
+                transformer_layer_spec = (
+                    get_transformer_block_with_experimental_attention_variant_spec(
+                        config=config, vp_stage=vp_stage
+                    )
+                )
+            elif args.num_experts:
                 assert not (config.transformer_impl == "inference_optimized")
                 # Define the decoder block spec
                 transformer_layer_spec = get_gpt_decoder_block_spec(

megatron/core/jit.py

@@ -7,12 +7,27 @@
 jit_fuser = torch.jit.script
 # nvFuser is deprecated in PyTorch JIT starting from 2.2
 
-try:
-    if is_torch_min_version("2.2.0a0"):
-        jit_fuser = torch.compile
-except ImportError:
 
-    def noop_decorator(func):
-        return func
+def noop_decorator(func):
+    '''No-op decorator'''
+    return func
 
+
+def enable_jit_fuser():
+    '''Enable the JIT fuser'''
+    global jit_fuser
+    try:
+        if is_torch_min_version("2.2.0a0"):
+            jit_fuser = torch.compile
+    except ImportError:
+
+        jit_fuser = noop_decorator
+
+
+def disable_jit_fuser():
+    '''Disable the JIT fuser'''
+    global jit_fuser
     jit_fuser = noop_decorator
+
+
+enable_jit_fuser()