[Feature] Add GLM-5.3-Flash F2: vision tower + projector (eager)

tests/model/test_moe_ignored_grad_sync.py

@@ -85,7 +85,7 @@ def _check_case(self, ep_size: int, sp_size: int) -> None:
         assert all(isinstance(placement, Replicate) for placement in parameter.placements)
         assert model.world_mesh is not None
         assert model.world_mesh.size() == self.world_size
-        assert parameter.device_mesh.size() == (ep_size if ep_size > 1 else self.world_size)
+        assert parameter.device_mesh == model.world_mesh
         assert any(parameter is p for _, p in model.trainable_parameters())
         with torch.no_grad():
             parameter.to_local().fill_(1.0)

xtuner/v1/data_proto/sequence_context.py

@@ -2,6 +2,7 @@
 from typing import cast
 
 import torch
+from torch.distributed._functional_collectives import all_gather_tensor_autograd
 from torch.distributed.device_mesh import DeviceMesh
 from typing_extensions import Self
 
@@ -461,8 +462,11 @@ def raw_inputs_embeds(self) -> torch.FloatTensor | None:
             return None
         if self.sequence_parallel_mesh is None or self.sequence_parallel_mesh.size() == 1:
             return self.inputs_embeds
-        gathered = gather_for_sequence_parallel(
-            self.inputs_embeds, dim=1, sp_group=self.sequence_parallel_mesh.get_group()
+        # MTP rolls these embeddings to construct its future-token input. Unlike
+        # integer input_ids, this gather must reduce-scatter gradients back to
+        # their owners, including gradients crossing an SP shard boundary.
+        gathered = all_gather_tensor_autograd(
+            self.inputs_embeds.contiguous(), gather_dim=1, group=self.sequence_parallel_mesh
         )
         self._raw_inputs_embeds = cast(torch.FloatTensor, gathered)
         return self._raw_inputs_embeds

xtuner/v1/model/base.py

@@ -680,23 +680,44 @@ def traverse(module):
 
                 for hf_name in hf_name_list:
                     if any(re.search(p, hf_name) for p in patterns):  # type: ignore
-                        if not isinstance(param, DTensor):
-                            dist_param = nn.Parameter(
-                                distribute_tensor(
-                                    param, self.world_mesh, [Replicate() for _ in range(self.world_mesh.ndim)]
-                                ),
-                                requires_grad=param.requires_grad,
-                            )
-                            module.register_parameter(name, dist_param)
-                            ignored_params.add(dist_param)
-                        else:
-                            # param is already a DTensor (e.g. distributed by
-                            # MoE._replicate_other_params on ep_mesh before _fully_shard
-                            # is called). We skip re-distributing on world_mesh and just
-                            # add it to ignored_params so FSDP leaves it alone.
-                            # ASSUMPTION: fp32 distribution always happens AFTER any
-                            # prior EP distribution, so the existing placement is correct.
-                            ignored_params.add(param)
+                        if isinstance(param, DTensor):
+                            # Case 1: Already on the target world mesh, e.g. an FP32
+                            # replica with [Replicate(), Replicate()] on an FSDP x EP
+                            # mesh. Preserve all existing placements on this mesh.
+                            if param.device_mesh == self.world_mesh:
+                                ignored_params.add(param)
+                                break
+
+                            # Case 2: On another mesh with a non-Replicate placement.
+                            # For Shard, to_local() returns only a parameter shard;
+                            # for Partial, it returns a value pending reduction.
+                            # Neither can be replicated as the full parameter.
+                            # Preserve these layouts; converting them is outside this fix.
+                            if any(not placement.is_replicate() for placement in param.placements):
+                                ignored_params.add(param)
+                                break
+
+                        # Case 3: A plain Tensor/Parameter becomes replicated across
+                        # every dimension of the target world mesh.
+                        #
+                        # Case 4: A fully replicated DTensor on another mesh, e.g.
+                        # the EP (or EP x expert-TP) submesh used by MoE initialization.
+                        # If that submesh covers fewer ranks, include the missing
+                        # FSDP rows so their FP32 gradients also synchronize.
+                        # If it already covers all ranks, normalize the mesh layout
+                        # without expanding the actual replica set.
+                        # Each rank holds the full parameter, so to_local().detach()
+                        # supplies a local leaf for rebuilding the DTensor. The
+                        # replacement remains ignored by FSDP.
+                        local_param = param.to_local().detach() if isinstance(param, DTensor) else param
+                        dist_param = nn.Parameter(
+                            distribute_tensor(
+                                local_param, self.world_mesh, [Replicate() for _ in range(self.world_mesh.ndim)]
+                            ),
+                            requires_grad=param.requires_grad,
+                        )
+                        module.register_parameter(name, dist_param)
+                        ignored_params.add(dist_param)
                         break
 
             for child in module.children():

xtuner/v1/model/compose/glm53/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from .glm53_config import Glm53ProjectorConfig, Glm53VisionConfig
+from .modeling_projector import Glm53Projector
+from .modeling_vision import Glm53VisionModel
+from .vision_utils import flatten_video_grid_thw
+
+
+__all__ = [
+    "Glm53VisionConfig",
+    "Glm53ProjectorConfig",
+    "Glm53VisionModel",
+    "Glm53Projector",
+    "flatten_video_grid_thw",
+]

xtuner/v1/model/compose/glm53/glm53_config.py

@@ -0,0 +1,72 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+"""GLM-5.3-Flash vision configs, see doc/xtuner_glm5p3flash_design.md F2
+§5.1/§3.7.4.
+
+Field names/defaults are checked against the real checkpoint's ``vision_config`` (transformers
+5.17.0, `zai-org/GLM-5.3-Flash`), not copied from the design doc's placeholder table:
+
+- the checkpoint field is ``num_heads``, not ``num_attention_heads`` (HF aliases the latter via
+  ``attribute_map``; XTuner has no such alias mechanism, so the field here is named to match the
+  checkpoint directly);
+- the checkpoint has no ``rope_parameters`` key at all -- HF's ``AutoConfig`` fills the default
+  ``{"rope_theta": 10000.0, "rope_type": "axial"}`` when absent, so that's the default here too
+  (declaring it required would block construction from the real config.json);
+- the checkpoint field is ``out_hidden_size``, not ``text_hidden_size``.
+"""
+
+from typing import Literal
+
+from pydantic import ConfigDict
+
+from xtuner.v1.model.base import XTunerBaseModelConfig
+
+
+class Glm53VisionConfig(XTunerBaseModelConfig):
+    model_config = ConfigDict(title="GLM-5.3-Flash vision config for xtuner", extra="forbid")
+
+    in_channels: int = 3
+    depth: int = 24
+    hidden_size: int = 1024
+    num_heads: int = 16
+    intermediate_size: int = 4096
+    patch_size: int = 14
+    temporal_patch_size: int = 2
+    spatial_merge_size: int = 2
+    rms_norm_eps: float = 1e-5
+    hidden_act: str = "silu"
+    swiglu_limit: float = 10.0
+    rope_parameters: dict = {"rope_theta": 10000.0, "rope_type": "axial"}
+    attention_bias: bool = True
+    attention_dropout: float = 0.0
+    attn_impl: Literal["flash_attention", "flex_attention", "eager_attention"] = "flash_attention"
+    fully_shard: bool = True
+
+    def build(self):
+        from .modeling_vision import Glm53VisionModel
+
+        return Glm53VisionModel(self)
+
+    @property
+    def hf_config(self):
+        return None
+
+
+class Glm53ProjectorConfig(XTunerBaseModelConfig):
+    model_config = ConfigDict(title="GLM-5.3-Flash projector config for xtuner", extra="forbid")
+
+    vision_hidden_size: int = 1024
+    out_hidden_size: int = 4096
+    spatial_merge_size: int = 2
+    projection_intermediate_size: int = 10240
+    hidden_act: str = "silu"
+    swiglu_limit: float = 10.0
+    fully_shard: bool = True
+
+    def build(self):
+        from .modeling_projector import Glm53Projector
+
+        return Glm53Projector(self)
+
+    @property
+    def hf_config(self):
+        return None

xtuner/v1/model/compose/glm53/modeling_projector.py

@@ -0,0 +1,116 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+"""GLM-5.3-Flash projector: `downsample` (2x2 spatial merge) + `merger`, see
+design doc F2 §5.2.
+
+Local naming mirrors HF exactly (`merger.proj/post_projection_norm/gate_proj/up_proj/down_proj`,
+`downsample`) per §6.1, so `from_hf()` needs no per-model copy loop -- only `to_hf_key_list`
+remaps the `model.visual.` prefix.
+"""
+
+import torch
+import torch.nn as nn
+from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy
+from typing_extensions import override
+
+from xtuner.v1.config import FSDPConfig
+from xtuner.v1.model import BaseModel, TorchCompileOption
+from xtuner.v1.ops.act_fn import get_act_fn
+from xtuner.v1.utils.init_weight import default_init_weights
+
+from .glm53_config import Glm53ProjectorConfig
+from .modeling_vision import init_world_mesh
+
+
+# The projector is one Conv2d downsample plus the merger's clamped-SwiGLU chain, with no
+# kernel dispatch or collective inside, so it takes a full graph.
+GLM53_PROJECTOR_COMPILE_CFG: dict[str, TorchCompileOption] = {
+    "xtuner.v1.model.compose.glm53.modeling_projector.Glm53Projector.forward": TorchCompileOption(
+        fullgraph=True
+    ),
+}
+
+
+class Glm53VisionPatchMerger(nn.Module):
+    def __init__(self, config: Glm53ProjectorConfig) -> None:
+        super().__init__()
+        dim, context_dim = config.out_hidden_size, config.projection_intermediate_size
+        self.proj = nn.Linear(dim, dim, bias=False)
+        self.post_projection_norm = nn.LayerNorm(dim)
+        self.gate_proj = nn.Linear(dim, context_dim, bias=False)
+        self.up_proj = nn.Linear(dim, context_dim, bias=False)
+        self.down_proj = nn.Linear(context_dim, dim, bias=False)
+        self.act1 = nn.GELU()
+        self.act_fn = get_act_fn(config.hidden_act)
+        self.swiglu_limit = config.swiglu_limit
+
+    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        hidden_state = self.proj(hidden_state)
+        hidden_state = self.act1(self.post_projection_norm(hidden_state))
+        gate = self.gate_proj(hidden_state).clamp(max=self.swiglu_limit)
+        up = self.up_proj(hidden_state).clamp(min=-self.swiglu_limit, max=self.swiglu_limit)
+        return self.down_proj(self.act_fn(gate) * up)
+
+
+class Glm53Projector(BaseModel):
+    config: Glm53ProjectorConfig
+
+    def __init__(self, config: Glm53ProjectorConfig) -> None:
+        super().__init__(config)  # type: ignore[arg-type]
+        self.spatial_merge_size = config.spatial_merge_size
+        self.downsample = nn.Conv2d(
+            in_channels=config.vision_hidden_size,
+            out_channels=config.out_hidden_size,
+            kernel_size=config.spatial_merge_size,
+            stride=config.spatial_merge_size,
+        )
+        self.merger = Glm53VisionPatchMerger(config)
+
+        self._hf_prefix = "model.visual."
+        self._init_load_spec()
+
+    @property
+    @override
+    def default_compile_cfg(self) -> dict[str, TorchCompileOption]:
+        return GLM53_PROJECTOR_COMPILE_CFG
+
+    def to_hf_key_list(self, key: str) -> list[str]:
+        return [self._hf_prefix + key]
+
+    @torch.no_grad()
+    def init_weights(self) -> None:
+        initialized = default_init_weights(self)
+        if missing := {name for name, _ in self.named_parameters()} - initialized:
+            raise RuntimeError(f"{missing} is not initialized")
+
+    @override
+    def fully_shard(self, fsdp_config: FSDPConfig):
+        self.fsdp_config = fsdp_config
+        self.fsdp_mesh = init_world_mesh()
+
+        if fsdp_config.requires_grad:
+            for module in self.modules():
+                for p_name, param in module.named_parameters(recurse=False):
+                    if param.requires_grad:
+                        setattr(module, p_name, torch.nn.Parameter(param.to(dtype=torch.float32)))
+        else:
+            for param in self.parameters():
+                param.requires_grad = False
+
+        if self.config.fully_shard:
+            mp_policy = MixedPrecisionPolicy(
+                param_dtype=fsdp_config.param_dtype, reduce_dtype=fsdp_config.reduce_dtype
+            )
+            self._fully_shard(
+                mesh=self.fsdp_mesh,
+                mp_policy=mp_policy,
+                reshard_after_forward=True,
+                offload_policy=CPUOffloadPolicy() if fsdp_config.cpu_offload else None,
+            )
+        return self
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        m = self.spatial_merge_size
+        hidden_states = hidden_states.view(-1, m, m, hidden_states.shape[-1])
+        hidden_states = hidden_states.permute(0, 3, 1, 2)
+        hidden_states = self.downsample(hidden_states).view(-1, self.config.out_hidden_size)
+        return self.merger(hidden_states)