Skip to content

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

Open
jayhenry wants to merge 5 commits into
feat/glm53flash-f1-vl-datafrom
feat/glm53flash-f2-vision-tower
Open

jayhenry wants to merge 5 commits into
feat/glm53flash-f1-vl-datafrom
feat/glm53flash-f2-vision-tower

Conversation

@jayhenry

@jayhenry jayhenry commented Sep 23, 2026 •

Copy link
Copy Markdown
Collaborator

Stack (bottom to top):

  1. [Feature] Add GLM-5.3-Flash F0: 25B cropped reference checkpoint builder #2105 feat/glm53flash-materialize-full-f0 → main
  2. [Feature] Add GLM-5.3-Flash F3: Kimi Delta Attention (KDA) #2106 feat/glm53flash-f3-kda → feat/glm53flash-materialize-full-f0
  3. [Feature] Add GLM-5.3-Flash F4: mHC four-stream residual #2107 feat/glm53flash-f4-mhc → feat/glm53flash-f3-kda
  4. [Feature] Add GLM-5.3-Flash F5: NoPE DSA + KPool indexer + clamped SwiGLU #2108 feat/glm53flash-f5-nope-dsa → feat/glm53flash-f4-mhc
  5. [Feature] Add GLM-5.3-Flash F1: VL data preprocessing pipeline #2109 feat/glm53flash-f1-vl-data → feat/glm53flash-f5-nope-dsa
  6. [Feature] Add GLM-5.3-Flash F2: vision tower + projector (eager) #2110 feat/glm53flash-f2-vision-tower → feat/glm53flash-f1-vl-data ← you are here
  7. [Feature] Add GLM-5.3-Flash F6 core: text model + MTP + compose model #2111 feat/glm53flash-f6-text-moe → feat/glm53flash-f2-vision-tower

Base is #2109's branch (layer 5). Review only this PR's own diff.


Summary

Stack layer 6/7 of GLM-5.3-Flash support (base: layer 5, F1 VL data pipeline). The F2 vision tower and projector, taken from "eager parity only" to something that survives production training.

F2 core

Native XTuner vision tower and projector (xtuner/v1/model/compose/glm53/), split per the design doc's §5.2 module boundary: HF's single Glm5NextVisionModel becomes vision_tower (patch_embed → 2D axial RoPE → 24 blocks → post_layernorm) plus multi_modal_projector (Conv2d downsample + clamped-SwiGLU merger). cu_seqlens/position_ids are reimplemented natively rather than imported from transformers.vision_utils, matching qwen3_vl's established pattern. The vision RMSNorm follows HF/Automodel's cast-then-affine rounding, which matters in bf16.

Vision sequence parallel (design doc §9.3/§9.6)

Previously the tower rejected sequence_parallel_mesh.size() > 1 outright — the attention had half an Ulysses path but nothing sharded the patch sequence, which silently computes the wrong attention. Now implemented properly:

  • _shard_patches_for_sequence_parallel pads the global patch sequence up to a multiple of sp_size * spatial_merge_size ** 2 and splits it contiguously. Merge alignment is a hard constraint: the projector's downsample treats each consecutive group of 4 patches as one 2×2 block, so a shard boundary inside a group would merge patches owned by two different ranks — alignment is also what lets each rank run the projector on its own shard (local projector). Equal length is what Ulysses all-to-all requires.
  • The padding is declared as an extra grid_thw row rather than appended silently, so it forms its own cu_seqlens segment: padding patches can never attend to, or be attended by, a real image.
  • Position ids are sharded to match the local patches (RoPE runs before the all-to-all); cu_seqlens stays global (after the all-to-all, attention sees the whole sequence).

Three production bugs

  1. attn_impl="flash_attention" was broken, not merely unverified. It died with Expected a value of type 'Tensor' for argument 'max_seqlen_q' but instead found type 'int' — the tower was the only attention call site materializing max_seqlen via int(...item()), which flash_attn_varlen_func v2 rejects and which also forces a device sync. Now passed as a tensor, like every other call site does with SequenceContext.max_length_q.
  2. pixel_values was never moved to the device. SequenceContext deliberately leaves the media tensors on CPU (documented there as "sharded by the model side, then moved"), so real training died on step 1 with Input type (CPUBFloat16Type) and weight type (CUDABFloat16Type) should be the same. Every unit test passed device tensors in, so nothing caught it. The move happens after the SP split, so only this rank's shard is transferred.
  3. The ViT blocks had no activation checkpointing — the actual cause of multimodal OOM. The tower runs over raw patches, and a packed VL sample carries far more of them than the LLM sequence has tokens (an 8192-token pack measured img_tokens: 14688; a 16384-token pack ~31k), so 24 blocks of retained attention + 4096-wide MLP activations dominated peak memory. VL SFT OOMed in backward at 16k pack and still died on step 2 at 8k, long before the language tower was near its limit. Now wired through apply_activation_checkpointing honouring vision_recompute_ratio/checkpoint_preserve_rng_state, the same knobs qwen3_vl's tower uses. Peak drops to 85.85 GB at 16k pack.

Compile config

default_compile_cfg was empty for both modules, so the vision half stayed eager even under MODEL_COMPILE=1. The ViT block takes fullgraph=False (it dispatches into an attention kernel, and under SP it contains collectives); the projector takes a full graph.

Test Plan

tests/model/test_glm53_vision.py 14/14:

  • Real-checkpoint weight mapping: 339 + 8 = 347 parameters, from_hf(strict=False) with no missing/unloaded keys.
  • Forward parity vs HF Glm5NextVisionModel at rtol=0, atol=0 (fp32 and bf16): single image, multi-tubelet video, multi-image batch, plus backward gradient flow.
  • TestGlm53VisionSequenceParallel::test_sp_tower_and_projector_match_non_sp (2 GPU): 20 patches over two images — deliberately not a multiple of sp_size * merge_unit, exercising the padding branch — tower + projector under SP, gathered and trimmed, matches the non-SP result.
  • test_unaligned_patch_count_is_rejected_under_sp: a patch count that would split a merge block raises instead of mis-merging.
  • test_flash_attention_matches_eager (GPU): the flash path runs and agrees with the HF-aligned eager baseline.
  • test_cpu_pixel_values_are_moved_to_device: CPU pixel_values with a device grid_thw, the exact shape SequenceContext produces.

Known gaps

Vision-side FSDP2 multi-GPU parity is still covered only by the end-to-end smoke, not a dedicated test_vision_fsdp_parity.

jayhenry and others added 5 commits September 24, 2026 20:22
Implements the eager-precision core of the F2 milestone from
doc/xtuner_glm5p3flash_design.md: the native XTuner vision tower and
projector for GLM-5.3-Flash, split per the design doc's module boundary
(xtuner/v1/model/compose/glm53/). Authority for the F2 contract is
duanyanhui's xtuner_glm53_flash_vl_vision_design.md §4-§6/§9-§13 and
xtuner_glm53_flash_vl_develop_stages.md §2/§4/§5/§6 (stages 1/3/4/5);
this delivery covers stages 1 and 3 (skeleton + weight mapping, eager
bitwise forward) in full and stages 4/5 (FSDP2/compile, Vision SP) as
wiring only, not production-grade verification -- see Known gaps below
and in doc/progress.md.

- glm53_config.py: Glm53VisionConfig/Glm53ProjectorConfig, fields
  checked against the real checkpoint's vision_config (num_heads not
  num_attention_heads, out_hidden_size not text_hidden_size,
  rope_parameters defaulted since the checkpoint JSON omits it).
- modeling_vision.py: Glm53VisionModel (patch_embed -> 2D axial RoPE ->
  24 blocks with RMSNorm+QK-norm+attention+clamped-SwiGLU MLP ->
  post_layernorm). cu_seqlens/position_ids are natively reimplemented
  rather than imported from transformers.vision_utils at runtime, per
  the "native implementation, not an HF training-time dependency"
  principle (§3.2) and matching qwen3_vl's own established pattern.
- modeling_projector.py: Glm53Projector (downsample Conv2d + merger).
  XTuner's tower/projector split moves downsample into the projector,
  unlike HF's single Glm5NextVisionModel -- confirmed intentional via
  the design doc's HF-key-mapping table (§6.1), not an HF module
  boundary.
- vision_utils.py: flatten_video_grid_thw.

Test plan: tests/model/test_glm53_vision.py (8 tests, CPU only). Real-
checkpoint weight mapping (gated, skips without GLM_5_3_FLASH_PATH):
vision_tower (339) + projector (8) = 347 parameters, matching the
design doc's independently-recorded real-checkpoint vision key count;
from_hf(strict=False) leaves no missing/unloaded keys, spot-checked
tensors are torch.equal to the raw safetensors. Forward parity vs HF
Glm5NextVisionModel (tiny synthetic config, eager attention, rtol=0
atol=0): single image, multi-tubelet video (grid=[2,4,4], exercises
multi-segment cu_seqlens and t-repeated position ids), multi-image
batch, plus a backward gradient-flow smoke check. Ground truth is HF
model.visual as a whole (§5.2): XTuner's post_layernorm output is
compared against HF's pre-downsample state (recomputed via HF's own
downsample module, not just pooler_output).

Known gaps (design doc's own stage split -- stages 4/5 need multi-GPU
coordination -- recorded explicitly, not silently skipped): FSDP2
wiring (fully_shard) mirrors qwen3_vl's structural pattern but has no
multi-GPU parity test; attention supports sequence_parallel_mesh with
Ulysses all-to-all on q/k/v but the full merge-aligned-padding +
local-projector Vision SP design (§9.3/§9.6) and its SP=2/4 parity
suite are not implemented; torch.compile config not wired; production
attention kernel (FlashAttention/FlexAttention) tolerance matrix (§11.7)
not verified, only eager_attention; HF save round-trip not verified.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
bf16 parity needs the HF/Automodel cast-then-affine norm. Ulysses without
merge-aligned padding computes the wrong attention, so size>1 now errors.

Co-authored-by: Cursor <cursoragent@cursor.com>
…h kernel path

Closes three F2 known gaps that together kept the vision tower off the production training
path.

**Vision sequence parallel.** The tower previously rejected `sequence_parallel_mesh.size()>1`
outright: Ulysses on an unsharded patch sequence computes the wrong attention, and nothing split
the patches. Implements the design doc's §9.3/§9.6 merge-aligned padding + local projector:
`_shard_patches_for_sequence_parallel` pads the global patch sequence up to a multiple of
`sp_size * spatial_merge_size ** 2` and splits it contiguously, so (a) every shard holds a whole
number of 2x2 merge blocks and can therefore run the projector on its own shard, and (b) all
ranks carry the same patch count, which Ulysses all-to-all requires. The padding is declared as
an extra `grid_thw` row instead of being appended silently, which keeps it a separate
`cu_seqlens` segment -- padding patches can never attend to, or be attended by, a real image.
Position ids are sharded to match the local patches (RoPE is applied before the all-to-all)
while `cu_seqlens` stays global (the all-to-all hands attention the whole sequence).

**Flash attention was broken, not just unverified.** `attn_impl="flash_attention"` died with
`Expected a value of type 'Tensor' for argument 'max_seqlen_q' but instead found type 'int'`:
the tower was the only attention call site materializing `max_seqlen` via `int(...item())`,
which `flash_attn_varlen_func` v2 rejects outright and which also forces a device sync. Now
passed as a tensor like every other call site does with `SequenceContext.max_length_q`.

**Compile config.** `default_compile_cfg` was empty for both modules, so the vision half stayed
eager even under `MODEL_COMPILE=1`. The ViT block takes `fullgraph=False` (it dispatches into an
attention kernel, and under SP it also contains collectives); the projector, which is a Conv2d
plus the merger's clamped-SwiGLU chain with no dispatch or collective, takes a full graph.

Test Plan:
  - `TestGlm53VisionSequenceParallel::test_sp_tower_and_projector_match_non_sp` (2 GPU): 20
    patches over two images -- deliberately not a multiple of `sp_size * merge_unit`, so the
    padding branch is exercised -- tower + projector under SP, gathered and trimmed, matches the
    non-SP result.
  - `test_unaligned_patch_count_is_rejected_under_sp`: a patch count that would split a merge
    block raises instead of silently mis-merging.
  - `TestGlm53VisionProductionAttention::test_flash_attention_matches_eager` (GPU): the flash
    path now runs and agrees with the HF-aligned eager baseline in bf16.
  - tests/model/test_glm53_vision.py 13/13.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…edding

`SequenceContext` deliberately does not move `pixel_values`/`pixel_values_videos` in its `.to()`
(documented there as "sharded by the model side, then moved"), so the tower receives CPU tensors
in real training and died on the very first step with `Input type (CPUBFloat16Type) and weight
type (CUDABFloat16Type) should be the same`. Every unit test passed device tensors in directly,
so nothing caught it.

The move happens after the SP split, not before, so only this rank's shard is transferred -- the
reason `SequenceContext` leaves the tensor alone in the first place. It follows the module's own
parameter device rather than the global one, so a CPU-resident tower (used by the parity tests)
keeps working.

Test Plan: `test_cpu_pixel_values_are_moved_to_device` feeds CPU `pixel_values` with a
device-resident `grid_thw` -- the exact shape `SequenceContext` produces -- to a CUDA tower.
tests/model/test_glm53_vision.py 14/14.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The vision tower never applied activation checkpointing, so all 24 ViT blocks retained their
activations. That is not a small oversight for this model: the tower runs over *raw* patches,
and a packed VL sample carries far more of them than the LLM sequence has tokens -- an 8192-token
pack measured `img_tokens: 14688` -- so the retained attention plus 4096-wide MLP activations of
24 blocks dominate the step's peak memory. End-to-end VL SFT on the F0 25B checkpoint OOMed in
backward at `PACK_MAX_LENGTH=16384`, and still OOMed on step 2 at 8192, before the language
tower was anywhere near its limit.

Wires the blocks through `apply_activation_checkpointing` honouring `vision_recompute_ratio`
(default 1.0) and `checkpoint_preserve_rng_state`, the same knobs qwen3_vl's tower already uses.

Test Plan: tests/model/test_glm53_vision.py 14/14 (the checkpointed path keeps HF bitwise parity
and the SP parity result); end-to-end VL SFT memory result recorded in doc/progress.md.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@jayhenry
jayhenry force-pushed the feat/glm53flash-f2-vision-tower branch from bf6a4aa to 26438cd Compare September 24, 2026 20:22

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants