Conversation
jayhenry
added this pull request to stack #2112
September 23, 2026 06:54
This was referenced Sep 23, 2026
jayhenry
force-pushed
the
feat/glm53flash-f2-vision-tower
branch
3 times, most recently
from
September 24, 2026 17:55
3f02459 to
bf69e1a
Compare
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
force-pushed
the
feat/glm53flash-f2-vision-tower
branch
from
September 24, 2026 20:22
bf6a4aa to
26438cd
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stack (bottom to top):
feat/glm53flash-materialize-full-f0→mainfeat/glm53flash-f3-kda→feat/glm53flash-materialize-full-f0feat/glm53flash-f4-mhc→feat/glm53flash-f3-kdafeat/glm53flash-f5-nope-dsa→feat/glm53flash-f4-mhcfeat/glm53flash-f1-vl-data→feat/glm53flash-f5-nope-dsafeat/glm53flash-f2-vision-tower→feat/glm53flash-f1-vl-data← you are herefeat/glm53flash-f6-text-moe→feat/glm53flash-f2-vision-towerSummary
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 singleGlm5NextVisionModelbecomesvision_tower(patch_embed → 2D axial RoPE → 24 blocks → post_layernorm) plusmulti_modal_projector(Conv2d downsample + clamped-SwiGLU merger).cu_seqlens/position_idsare reimplemented natively rather than imported fromtransformers.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() > 1outright — 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_parallelpads the global patch sequence up to a multiple ofsp_size * spatial_merge_size ** 2and splits it contiguously. Merge alignment is a hard constraint: the projector'sdownsampletreats 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.grid_thwrow rather than appended silently, so it forms its owncu_seqlenssegment: padding patches can never attend to, or be attended by, a real image.cu_seqlensstays global (after the all-to-all, attention sees the whole sequence).Three production bugs
attn_impl="flash_attention"was broken, not merely unverified. It died withExpected a value of type 'Tensor' for argument 'max_seqlen_q' but instead found type 'int'— the tower was the only attention call site materializingmax_seqlenviaint(...item()), whichflash_attn_varlen_funcv2 rejects and which also forces a device sync. Now passed as a tensor, like every other call site does withSequenceContext.max_length_q.pixel_valueswas never moved to the device.SequenceContextdeliberately leaves the media tensors on CPU (documented there as "sharded by the model side, then moved"), so real training died on step 1 withInput 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.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 throughapply_activation_checkpointinghonouringvision_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_cfgwas empty for both modules, so the vision half stayed eager even underMODEL_COMPILE=1. The ViT block takesfullgraph=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:
from_hf(strict=False)with no missing/unloaded keys.Glm5NextVisionModelatrtol=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 ofsp_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: CPUpixel_valueswith a devicegrid_thw, the exact shapeSequenceContextproduces.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.