Add Qwen3-VL support with Megatron-FSDP

megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py

@@ -1005,9 +1005,11 @@ def _register_pre_backward_param_unshard_hook(module):
                             with_kwargs=True,
                         )
                     )
-                grad_acc_param_list = list(module.parameters())
+                grad_acc_param_list = [p for p in module.parameters() if p.requires_grad]
             else:
-                grad_acc_param_list = list(module.parameters(recurse=False))
+                grad_acc_param_list = [
+                    p for p in module.parameters(recurse=False) if p.requires_grad
+                ]
 
             for param in grad_acc_param_list:
                 self.grad_acc_hooks[f"grad_acc and reduce for {self.param_to_name[param]}"] = (

megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py

@@ -2473,6 +2473,8 @@ def _reset_parameters(self, old_params, new_params):
             self.param_to_direct_module[new_param] = self.param_to_direct_module[old_param]
             del self.param_to_direct_module[old_param]
 
+            new_param.requires_grad_(old_param.requires_grad)
+
             for tp_attr in ["_mcore_tp", "_tp_partition_dim", "_tp_duplicated"]:
                 if getattr(old_param, tp_attr, None) is not None:
                     setattr(new_param, tp_attr, getattr(old_param, tp_attr))

megatron/core/models/gpt/gpt_model.py

@@ -444,7 +444,7 @@ def _preprocess(
             # return this extra tensor
             # this is for backwards compatibility with
             # legacy unit tests, which break if you
-            # return a 6 tuple instead of 5.
+            # return a 7 tuple instead of 6.
             preproc_output += (rotary_pos_cos_sin,)
 
         return preproc_output

megatron/core/transformer/fsdp_dtensor_checkpoint.py

@@ -45,6 +45,7 @@
 from megatron.core import parallel_state
 from megatron.core.tensor_parallel.layers import copy_tensor_model_parallel_attributes
 from megatron.core.transformer.transformer_layer import TransformerLayer
+from megatron.core.utils import get_attr_wrapped_model
 
 
 def get_ep_layer_offset(num_experts: int | None = None) -> int:
@@ -183,7 +184,10 @@ def handle_swiglu_in_state_dict(model, model_state_dict, optimizer_state_dict):
     assert HAVE_MEGATRON_FSDP, "This function requires Megatron-FSDP to be installed."
 
     # Extract num_experts from model config for expert parameter processing
-    num_experts = model.config.num_moe_experts if hasattr(model, 'config') else None
+    model_config = get_attr_wrapped_model(model, "config", allow_none=True)
+    num_experts = (
+        getattr(model_config, 'num_moe_experts', None) if model_config is not None else None
+    )
 
     def intersection(s1, s2):
         # Only works for step=1