Qwen3.5 support series

This is part of a 5-PR series adding Qwen3.5-VL support, split for review clarity.

Dev PRs (this series):

• [1/5] MTP packed-seq CP+THD fix — [Dev] fix(mtp): use padded cu_seqlens in MTP roll for THD with CP #4494
• [2/5] FSDP DTensor Bridge checkpoint compatibility — [dev] [2/5] Qwen3.5 support: FSDP DTensor Bridge checkpoint compatibility #4748 ← this PR
• [3/5] SharedExpertMLP meta init — [dev] [3/5] Qwen3.5 support: SharedExpertMLP meta init #4749
• [4/5] Interleaved MRoPE layout — [dev] [4/5] Qwen3.5 support: Interleaved MRoPE layout #4750
• [5/5] Qwen3.5-VL training example — [dev] [5/5] Qwen3.5 support: Qwen3.5-VL training example #4751

Main PRs (corresponding mirrors):

• [1/5] fix(mtp): use padded cu_seqlens in MTP roll for THD with CP #4495
• [2/5] [main] [2/5] Qwen3.5 support: FSDP DTensor Bridge checkpoint compatibility #4753
• [3/5] [main] [3/5] Qwen3.5 support: SharedExpertMLP meta init #4754
• [4/5] [main] [4/5] Qwen3.5 support: Interleaved MRoPE layout #4755
• [5/5] [main] [5/5] Qwen3.5 support: Qwen3.5-VL training example #4756

Summary

Three related fixes that together let fsdp_dtensor resume from a checkpoint converted with Megatron-Bridge:

1. split_swiglu_linear_fc1 / split_gdn_fused: accept plain Tensor (not only DTensor) for data.to_local(). Bridge writes some already-local tensors that previously crashed the splitter.

2. Stash missing top-level keys in _load_base_checkpoint: keys not present as DCP storage metadata (e.g. args, iteration) are stashed before dcp.load and restored after. Megatron-saved checkpoints contain these as BytesStorageMetadata; Bridge-converted ones don't, and DCP errors without this stash.

3. DistributedOptimizer._copy_model_params_to_main_params: under use_megatron_fsdp, return instead of raising. DCP already loads directly into the fp32 main buffer, and a post-load hook copies main → model weights — both are correct at this call site.

4. split_swiglu_linear_fc1 / split_gdn_fused use dist_param coordinates (cherry-pick of [DEV] fix(megatron-fsdp): compute SWiGLU/GDN split in item coordinates for non-DTensor optimizer states #4424 by @xuwchen): when data is a plain Tensor (FusedAdam optimizer state), data.numel() / data.shape describe the local item layout, not the global TP-sharded layout, so the splitter computed wrong data_size, view_shape, and per_tp_rank_shape. Switch to global_shape = dist_param.shape and dist_param.numel() everywhere, with a DTensor-shape consistency assert.

5. resolve conflict from [dev] Fix GDN DTensor splitting for FSDP checkpointing #4799

Also wires handle_gdn_in_state_dict into preprocess_fsdp_dtensor_state_dict.

Why bundled

All three fixes are needed together for the Bridge → fsdp_dtensor resume path to work end-to-end. Splitting further yields PRs that individually leave the path broken.

Note on key-unwrapping

An earlier revision of this branch also rewrote model.get_parameter(f"module.{key}") → model.get_parameter(key) (and the analogous optimizer-state lookups) on the assumption that unwrap_model strips MegatronFSDPModule. That root cause was resolved in #4393 ("Revert 'fix mfsdp unwrap stuck at MegatronFSDP [dev] (#4273)'"), which removed MegatronFSDPModule from unwrap_model's tuple, so the module. prefix is correct again. Those four lookup edits have been dropped from this PR.

Test plan

• Save a checkpoint under FSDP-DTensor with Megatron-LM, resume — unchanged.
• Convert a checkpoint with Megatron-Bridge, resume under FSDP-DTensor — succeeds (was previously failing on missing-key error from DCP, then on the GDN splitter for TP-local tensors).
• Hybrid model with GDN + SwiGLU MoE: loss curve matches pre-resume curve.

Notes

Extracted from a larger Qwen3.5-VL development branch to keep these fixes reviewable on their own.

🤖 Generated with Claude Code

Update (2026-05-14): model→main copy reverted

A previous revision of this PR (commit 4c2c3cbfb on dev / f4ba23b11 on main) implemented an FSDP copy_model_weights_to_main_weights and called it from _copy_model_params_to_main_params, on the theory that the silent return left fp32 main params uninitialized under --no-load-optim / --finetune.

Empirical buffer-diff probe showed that's wrong: under fsdp_dtensor, main_weight_buffer is exposed as the module's nn.Parameters via MegatronFSDP's state_dict pre-hook (_replace_param_with_distributed_if_needed), so the preceding model.load_state_dict(state_dict["model"]) writes the checkpoint's fp32 tensors directly into main_weight_buffer — full fp32 precision, even with --no-load-optim, because main params live in the model state dict (not the optimizer state dict) under fsdp_dtensor. The load_state_dict post-hook (copy_main_weights_to_model_weights) then casts main → bf16 to refresh model_weight_buffer.

A 50-step finetune+resume on a Qwen3.5-VL 0.8B model with and without the copy gave indistinguishable loss curves (wandb runs q8d7l7nm / o4tyuf5z / wwuz7y4z under https://wandb.ai/wplf/fsdp-fix-verify/). The pre-copy buffer diff was exactly bf16 rounding noise; the post-copy diff dropped to 0 only because main had been downgraded to bf16 precision.

The model→main copy commit has been reverted. The plain return is correct, and the comment in _copy_model_params_to_main_params now documents the mechanism so this isn't re-litigated.

Update (2026-05-14): TP-local GDN fast-path dropped

The original 2nd bullet (TP-local GDN split via split_dtensor) was the same fix proposed by @conver334 in #4799 against dev directly. We've dropped that bullet here so this PR doesn't textually conflict when #4799 merges; the remaining bullets are still complementary fixes that #4799 doesn't cover.