Repository navigation
[Fix] GLM-5.3-Flash F6.2: KDA correctness and FLA sequence parallelism - #2153
Merged
jayhenry merged 59 commits intoOct 13, 2026
Conversation
GLM-5.3-Flash introduced three copies of "unwrap a parameter before handing it to a kernel that does not understand DTensor", in two different flavours: mhc.unshard_hc_params used full_tensor(), while kda._to_local and nope_dsa_mla._to_local used to_local(). The two differ only when the placement is Shard -- to_local() silently returns this rank's slice, full_tensor() inserts an all-gather -- and neither is safe there, because with grad_placements=None both label the (per-rank different) gradient with the parameter's own placement and communicate nothing on the way back. xtuner/v1/utils/dtensor.py::materialize_full replaces all three: it accepts a plain tensor or a Replicate DTensor and raises on Shard, naming the parameter. Also folds the DTensor-aware LayerNorm that glm52 and glm53 had each defined into xtuner/v1/module/rms_norm/layer_norm.py; it is a generic norm, not a per-model one. Test Plan: tests/utils/test_dtensor.py::TestMaterializeFull covers pass-through, Replicate unwrap, and the Shard rejection. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Includes the GLM-5.3-Flash design assessment (doc/xtuner_glm5p3flash_design.{md,py}, doc/glm5p3flash_vs_glm5p2.md, demo/glm5p3_flash_demo.py) that this and every later F* milestone implements against.
Crops the published FP8 62-shard GLM-5.3-Flash checkpoint to a 5-main-layer
+ original-MTP-layer BF16 checkpoint (~24.9B params) for local validation,
following doc/xtuner_glm5p3flash_design.md F0. Dequantizes 128x128
block-scaled FP8 weights to BF16 and rewrites the nested text_config
schedule fields (layer_types/mlp_layer_types/indexer_types/kda_layers/
full_attn_layers).
Verified against the real checkpoint: 24.95B total params (matches the
design doc's 24.9B target), config loads under pinned transformers==5.17.0,
and Glm5NextForConditionalGeneration.from_pretrained runs a real forward
pass with finite loss. Findings recorded in doc/progress.md, including that
this transformers build's Glm5NextForConditionalGeneration has no MTP
forward path, so the renamed layers.5.* (MTP) weights load as harmless
UNEXPECTED keys.
The official GLM-5.3-Flash release also ships a native BF16 checkpoint
(120 shards, 45 layers, ~313B params, no quantization_config), separate
from the FP8 62-shard release this script was originally built against.
Investigated what would need to change to make the crop script and model
loading compatible with it.
Findings, none of which required a code change:
- The crop script's FP8 dequantization is already conditional per tensor
(triggers only when a `*_scale_inv` companion key exists), so it
transparently passes a native-BF16 source through untouched -- confirmed
via --dry-run (same 3100 tensors / 14 source shards selected) and a real
crop run.
- MoE expert weights are stored per-expert (`mlp.experts.{i}.
{gate,up,down}_proj.weight`, i=0..287) in *both* the FP8 and native BF16
releases -- this is the official on-disk layout, not something
introduced by dequantization. XTuner's own loader already expects
exactly this layout: Glm53TextMoE.to_hf_key_list() expands the internal
fused fused_w1w3.weight/fused_w2.weight into 288 per-expert HF keys for
both load and save, and hf_tensor_to_canonical() reassembles the
per-expert tensors read back into the fused shape. No per-expert-to-fused
conversion step was missing; loading only ever supports this one format,
matching both releases.
Used the native BF16 release to cross-validate F0's FP8 dequant math:
cropped 25B from both sources and diffed tensor-by-tensor. Tensors in
modules_to_not_convert (embed_tokens, layernorm, hc_* -- never quantized)
match bit-for-bit; FP8-quantized tensors (gate_proj, expert weights,
q_a_proj) match bit-for-bit on 92-96% of elements, with the remainder
differing by max_abs_diff ~2.3e-5 / mean_abs_diff ~1e-7-1e-6 -- consistent
with FP8 e4m3's quantization granularity, not a dequant bug.
Since the native-BF16-sourced crop is strictly better (no FP8 round-trip)
and both sources produce an identical on-disk layout, switched the
canonical `~/model/GLM-5.3-Flash-25B` (the path every GLM_5_3_FLASH_PATH-
gated test and sft_glm53_tiny.sh default to) to the native-BF16-sourced
crop. The original FP8-dequantized crop is kept at
`~/model/GLM-5.3-Flash-25B-fp8dequant` as the cross-validation baseline,
no longer the default.
tests/model/test_glm53_25b_crop.py: added
test_glm53_25b_crop_passes_through_native_bf16_weights_unchanged, covering
the no-quantization_config / no-*_scale_inv input path that was previously
untested (the existing fixture always simulated an FP8-quantized source).
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Adds xtuner/v1/module/attention/kda.py (KDAConfig / KimiDeltaAttention) following doc/xtuner_glm5p3flash_design.md F3, with module structure and parameter names matching the published checkpoint layout (separate q/k/v_conv1d, low-rank f_a_proj->f_b_proj forget gate with fp32 A_log/dt_bias, low-rank g_a_proj->g_b_proj output gate). Critically, the chunk_kda/fused_recurrent_kda call convention differs from the ~/github/xtuner-ncp-k3 reference: the installed fla 0.4.2's chunk_kda does not accept A_log/dt_bias/use_beta_sigmoid_in_kernel, so those get silently swallowed into **kwargs without raising if passed the ncp-k3 way, producing a KDA that trains but computes the wrong gate. The gate is now computed externally via fla.ops.kda.gate.fused_kda_gate before calling the kernel, matching HF's own Glm5NextTextLinearAttention reference forward. A regression test asserts these three kwargs stay absent from chunk_kda's signature. Relaxes TransformerConfig.linear_attention to GatedDeltaNetConfig | KDAConfig | None and the corresponding attention_config type annotations in dense/moe decoder layers and model builders (mypy-strict plumbing only, no behavior change for existing GatedDeltaNet layers). Verified with tests/model/test_glm53_kda.py (5 tests, single/2-GPU): gate parity against fla's naive reference, module-level parity against HF's Glm5NextTextLinearAttention (weights bridged per the published<->HF key mapping), packed multi-document isolation, and Ulysses SP parity. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Compare the chunked KDA forward output and every module gradient against HF's Glm5NextTextLinearAttention on GPU, including the split causal convolutions that HF fuses into one conv1d.
Port the dispatch fix and its regression tests onto F3. Preserve the existing F3 HF backward and signature guards; explicitly initialize raw gate parameters in standalone test fixtures. Original-Commit: d3e2d30
Include the causal-convolution custom-op dependency absent from F3 and route KDAShortConvolution through it. Forward retains activation-dtype rounding while FP32 parameters retain FP32 dw/db. The GLM-5.3 F6 fp32_keys_pattern change is intentionally excluded because that model configuration does not exist in F3; callers must retain convolution master parameters in FP32. Original-Commit: 974c39f
Include the fused-gate custom-op dependency absent from F3 and wire it into the KDA module. Sum dt_bias gradients in FP32 before restoring projection-gradient dtype; keep the established FLA forward formula. Original-Commit: 83d247b
Include the chunked-KDA custom-op dependency absent from F3 and wire it into the module. Use a fixed-layout FP32 gate prefix with FLA forward/backward stages. Preserve the existing upstream FLA gate-signature guard and add packed/head-slice, recurrent-reference, and fullgraph regressions. Original-Commit: 586b43f
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Eval-mode checkpoint replay can still switch kernels, and unsupported tensor arguments are handled incorrectly.
3 open findings
What changed in this PR
Improves KDA training correctness, gradient precision, checkpoint dispatch, and Ulysses stability.
Changes:
- Adds compile-friendly KDA gate, convolution, chunked kernel, and fixed-prefix operators.
- Uses chunked KDA whenever training or gradients are enabled.
- Adds precision, compilation, dispatch, checkpointing, packed-sequence, and SP tests.
| File | Description |
|---|---|
xtuner/v1/ops/kda/fused_kda_gate.py |
Adds FP32 gate-gradient wrapper. |
xtuner/v1/ops/kda/chunk_kda.py |
Adds compile-friendly chunked KDA. |
xtuner/v1/ops/kda/chunk_cumsum.py |
Adds fixed-layout prefix scan. |
xtuner/v1/ops/kda/causal_conv1d.py |
Preserves FP32 convolution gradients. |
xtuner/v1/ops/kda/__init__.py |
Initializes the KDA operators package. |
xtuner/v1/module/attention/kda.py |
Integrates operators and revises dispatch. |
tests/ops/test_kda_gate_precision.py |
Tests gate precision and compilation. |
tests/ops/test_kda_forward_precision.py |
Tests prefix, SP, and recurrent parity. |
tests/ops/test_kda_conv_precision.py |
Tests convolution gradient precision. |
tests/model/test_glm53_kda.py |
Tests dispatch, parity, and checkpointing. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Keep all heads on each rank's local sequence shard and exchange convolution history and recurrent boundary states instead of redistributing full projected tensors through Ulysses. Use the existing XTuner SP mesh without adding a parallelism dimension. Make FLA sequence SP the default in KDAConfig and KimiDeltaAttention; retain explicit sp_impl="ulysses" for the original path. Reuse the tested local FLA 0.4.2 backward adapter with corrected document/head gate offsets. Preserve F31 precision operators, keep L2-normalization backward inside the custom op, and use the existing gated-norm wrapper for fullgraph compilation. Installed dependencies and the conda environment remain unchanged. FLA SP requires equal contiguous shards with at least convolution-width-minus-one tokens per rank. Performance on this F31 branch: one KDA layer, eight H200 GPUs, SP8, batch 1, 131072 global tokens, hidden size 4096, 64 heads x 128, convolution width 4, BF16 activations/ordinary parameters with FP32 convolution and gate masters. Three warmups and ten timed iterations; medians of each iteration's maximum rank: - Ulysses: 26.23 ms forward, 49.79 ms backward, 76.00 ms total. - FLA sequence SP: 12.90 ms forward, 31.48 ms backward, 44.35 ms total. - FLA default with fullgraph compilation: 13.00 ms forward, 31.39 ms backward, 44.27 ms total. Eager FLA is 1.71x faster, reducing total time by 41.6%. Peak allocated memory is 10.125 GB/GPU for Ulysses, 10.129 GB for eager FLA, and 10.398 GB for compiled FLA. Timings include all KDA projections, convolution, recurrent core, normalization and SP activation-gradient communication; exclude MoE, mHC, outer checkpointing, optimizer/loss and outer parameter-gradient synchronization. Cold compilation/autotuning is excluded. This is a synthetic KDA-only benchmark, not full-model training throughput. Validation: existing 35 KDA tests plus three SP8 tests passed; reran four configuration/default/round-trip checks after changing the default. Fresh full-dimension 128K SP1 versus SP8 comparisons pass the existing 2% relative-L2 and 4% normalized-peak criteria for all outputs, input gradients, and all 15 parameter gradients on every rank, with both FP32 and native-dtype parameter-gradient summation. Input-gradient L2 is 0.435% eager and 0.446% compiled; worst parameter L2 with FP32 summation is 0.374% in both modes. Fullgraph forward/backward and the default config build are validated at 128K/SP8. Ruff and whitespace checks pass.
JT-Ushio
changed the title
GLM-5.3 HF parity tests require the glm5_next API introduced in Transformers 5.17. The upgrade also changes Qwen video interpolation baselines. Use the public video processor API, preserve fp32 vision RoPE, and mirror HF gated norm rounding in parity mode.
Adds xtuner/v1/module/decoder_layer/mhc.py (MHCConfig / hc_split_sinkhorn / hc_pre / hc_post / unshard_hc_params) and xtuner/v1/ops/hc_post.py (fused Triton hc_post_fused), following doc/xtuner_glm5p3flash_design.md F4. The core math is ported from xtuner's dsv4 branch (DeepSeek-V4's Hyper- Connections, commit 01c31a8, not merged into this branch) into the model-agnostic public location the design doc calls for, dropping V4's XTUNER_V4_HF_PARITY global toggle (the default bf16-fast path already degrades to HF-exact math under fp32 inputs, which the tests use as the parity anchor) and the unported 721-line TileKernels backend (left as a documented NotImplementedError gap rather than a blind port with no hardware to validate it against). Adds xtuner/v1/model/moe/glm53/decoder_layer.py: Glm53DenseDecoderLayer (overrides DenseDecoderLayer._forward) and Glm53MoEDecoderLayer, which only overrides MoEDecoderLayer's _pre_moe_forward/_post_moe_forward seams so the ~400-line EP/dispatcher/domino-micro-batch pipeline stays untouched; the mHC residual rides through those methods as an opaque _MHCResidual payload in place of the base class's plain Tensor. mhc_cfg=None degrades both layers to the base class's ordinary residual math unchanged, for reuse by the MTP layer in F6. Also fixes a real bug surfaced while testing: KDA's dt_bias/A_log were left torch.empty-uninitialized in xtuner/v1/module/attention/kda.py; now initialized to match HF's Glm5NextTextForgetGate _init_weights policy (A_log zeroed when a safe gate lower bound is set, dt_bias log-uniform). Verified with 26 tests across tests/model/test_glm53_mhc.py (Sinkhorn math + exact HF parity in fp32, bf16 parity within ULP tolerance), tests/ops/test_hc_post.py (fused kernel vs eager reference, forward/ backward/compile), and tests/model/test_glm53_decoder_layer.py (both decoder layers, mhc_cfg=None passthrough equivalence, gradient flow). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
HF normalizes the flattened mHC streams with the text-config RMS epsilon, not the hyper-connection eps. The hc_pre norm_eps parameter existed but all four call sites relied on the 1e-6 default, so checkpoint gradients drifted from HF. Pass input/post-attention layernorm variance_epsilon at each site; the five-layer crop gradient oracle now agrees with HF within 5%.
The build-only MoE test compiled shared decoder class methods globally and made later GLM-5.3 cases fail by test order. Disable compilation for that constructor check. Move the dynamic cu_seqlens compile case to F6, where the KDA custom-op boundary is implemented.
…iGLU
Implements the F5 milestone from doc/xtuner_glm5p3flash_design.md: the
NoPE (qk_rope_head_dim=0) DeepSeek Sparse Attention layers, their KPool
indexer (pools of index_kpool consecutive tokens scored together instead
of per-token top-k), a new flash_mla_cudnn SparseMLA backend (FlashMLA
forward + existing cuDNN backward), and the clamped SwiGLU activation
used by GLM-5.3-Flash's dense/shared/MoE MLPs.
- xtuner/v1/ops/act_fn.py: native_clamped_swiglu + MoEActFnConfig support.
- xtuner/v1/module/decoder_layer/{dense,moe}_decoder_layer.py: wire
swiglu_limit through DenseMLP/MoEMLP.
- xtuner/v1/ops/sparse_mla/kpool.py: pool layout, causal ranges, and
top-k pool selection (torch reference + TileLang-backed production
path, reusing the existing indexer kernel unmodified since relu's
homogeneity makes head_dim^-0.5 movable from the relu argument into
the per-head weight without changing the result).
- xtuner/v1/ops/sparse_mla/flash_mla_cudnn.py: new SparseMLA backend.
- xtuner/v1/ops/sparse_mla/tilelang.py: widen the hardcoded 576 head-dim
check to a (head_dim, value_dim) whitelist.
- xtuner/v1/model/moe/glm52/dsa_mla.py: reject flash_mla_cudnn as an
indexer_backend (SparseMLABackend widened, but this backend has no
indexer counterpart).
- xtuner/v1/model/moe/glm53/nope_dsa_mla.py: NoPEDSAMLAConfig,
KPoolIndexer, NoPEDSAMultiLatentAttention.
- xtuner/v1/ops/sparse_mla/protocol.py: KPoolIndexerBackend (restricts
GLM-5.3-Flash's KPool indexer_backend to "torch"/"tilelang" -- unlike
GLM-5.2's 6-way DSAIndexerBackend, there's no cudnn_dsa/flash_mla/
deep_gemm_fp8/cute_dsl KPool kernel) and KPoolTopKIndicesProtocol.
- xtuner/v1/ops/sparse_mla/__init__.py: get_kpool_topk_indices(backend),
mirroring get_dsa_topk_indices's style -- explicitly raises for any
backend other than "torch"/"tilelang" instead of silently falling
through. KPoolIndexer resolves this once in __init__ into
self._topk_indices_fn instead of re-branching on every forward call.
- xtuner/v1/model/moe/glm53/nope_dsa_mla.py: NoPEDSAMultiLatentAttention.
forward branches on freeze_dsa_indexer (torch.no_grad() only when
frozen), mirroring GLM-5.2's per-token indexer, ahead of a future
differentiable indexer output -- today's kpool_topk_indices/
torch_kpool_topk_indices still only ever return an int32 index
tensor, so this doesn't yet change what's trainable (confirmed via a
freeze_dsa_indexer=False smoke run: indexer params get
requires_grad=True but no actual gradient). Calls self.indexer(...)
directly instead of through reuse_during_recompute, which retained
topk_ids' activations across the backward recompute pass to avoid
recomputing them; dropped since nothing currently offloads or reuses
that retained memory.
Known gaps recorded in doc/progress.md: tilelang sparse_mla_backend and
deep_gemm_fp8 indexer_backend are not implemented for NoPE (explicit
NotImplementedError, not silent fallback); KPool/NoPE-DSA sequence-
parallel path is written per design but not yet GPU-tested at sp_size>1;
whether index topk ids can be offloaded is still an open question, not
attempted here.
Test plan: tests/model/test_glm53_dsa.py (12), test_glm53_nope_dsa_mla.py
(2), test_flash_mla_cudnn_sparse_mla.py (3) all pass; GLM-5.2's existing
tests/module/attention/test_dsa_mla.py (15) rerun clean against the
shared tilelang.py/dsa_mla.py edits, confirming no regression.
Numerical oracle: transformers 5.17.0's glm5_next model. Two near-tied
top-k floating-point sensitivities were root-caused via seed sweeps and
a targeted trace script (not worked around by loosening tolerances
blindly) -- see doc/progress.md F5 section for the full analysis.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Also folds in what were originally two follow-up commits, now part of F5 itself
rather than separate history:
- Give the DSA backends explicit defaults instead of inheriting: both
DSAMLAConfig/NoPEDSAMLAConfig fields now default to tilelang directly,
with indexer_backend never falling back to sparse_mla_backend (the two
name different, only partially overlapping backend vocabularies).
- Bound the KPool indexer's logits tile with query chunking, mirroring the
existing tilelang DSA selector's query-chunk bound.
…2143) * [Fix] Define and forward the TileLang DSA indexer selector 3de92e7 introduced the DeepSelect selector switch but referenced `selector` inside tilelang_indexer_topk_from_ranges without declaring it in the signature, so every call raised NameError. GLM-5.3's default KPool indexer backend (indexer_backend="tilelang") routes through this wrapper, so training could not start a single step on the f5 stack. The outer tilelang_dsa_topk_indices also accepted a selector but dropped it when calling the wrapper, so the tilelang_deepselect backend silently ran the torch.topk kernel instead of DeepSelect's radix-select kernel. - Declare selector in tilelang_indexer_topk_from_ranges, defaulting to "torch" to preserve the pre-DeepSelect behavior for existing callers (the KPool indexer). - Forward selector from tilelang_dsa_topk_indices. - Add CPU routing tests for the wrapper plus a CUDA-gated forwarding test that would have caught both regressions. * [Fix] Guard NoPE DSA config against unsupported settings at config time Two config-time guards for NoPEDSAMLAConfig, closing review findings #1 and InternLM#3 from the GLM-5.3-Flash stack review: - freeze_dsa_indexer=False now raises ValueError in build(), mirroring GLM-5.2's guard. The indexer only returns int32 indices and its gradients are cut by no_grad, so the flag silently left indexer params requires_grad=True -- paying optimizer state and activation memory for parameters that never train, with no error. - sparse_mla_backend='flash_mla_cudnn' (the production default) now requires num_attention_heads % 64 == 0 in the model validator, instead of failing at the first forward inside the FlashMLA kernel. The production config (64 heads) is unaffected; the torch reference backend keeps no alignment requirement. Also adds three tests covering both rejections and the torch-backend exemption. --------- Co-authored-by: liutong <liutong@pjlab.org.cn>
…#2147) b03bf01/fe0037f0 switched _cudnn_dsa_sparse_mla_backward_op from the log2-space LSE contract to natural-log LSE, but this backend kept forwarding the log2 value, so every GLM-5.3 training silently fed the cuDNN backward a 1.4427x-inflated LSE. Forward is unaffected; backward gradients come out systematically shrunk (grad_norm ~20-30% low) and loss drifts (+0.043 over 300 steps vs the pre-change baseline). Save and pass the natural-log softmax_lse instead. Co-authored-by: liutong <liutong@pjlab.org.cn>
The tiny GLM-5.2 cases selected TileLang dimensions and sparse MLA kernels outside their supported shape set. Keep the actual indexer path while using a supported index dimension and the torch sparse MLA reference. The F5 routed activation check now uses the F5 public config rather than importing the F6 text model.
The colocate test inherited WORLD_SIZE from other cases and requested 32 workers on an 8-GPU node. Direct sampled KL is noisy for greedy rollouts, so retain its finite check and bound the stable K3 metric. Check temporary CUDA tensor lifetime directly instead of requiring allocator reserved bytes to decrease.
Implements the F1 milestone from doc/xtuner_glm5p3flash_design.md: the
chat template, TokenizeFn (cache/runtime dual path), DataItem, and
Collator/SequenceContext plumbing for GLM-5.3-Flash's image/video SFT
data pipeline. Authority for the F1 contract is duanyanhui's
xtuner_glm53_flash_vl_vision_design.md §7-8 and
xtuner_glm53_flash_vl_develop_stages.md §3 (stage 2).
- xtuner/v1/data_proto/messages/glm53_chat.py: render_glm53_chat /
Glm53ChatMessages, adapted from the already-landed glm52_chat.py at
exactly its two documented diffs (reasoning_effort domain
{high,max}->{low,high,max}; media renders as the two-stage
<|begin_of_image|><|image|><|end_of_image|>-style placeholder marker
instead of GLM-5.2's <reminder> text). Verified byte-for-byte against
the real chat_template.jinja via tokenizer.apply_chat_template.
- xtuner/v1/datasets/mllm_tokenize_fn/glm53_vl_tokenize_fn.py:
Glm53VLTokenizeFunction/Config. Placeholder expansion and
mm_token_type_ids reuse the real Glm5NextProcessor's own public
replace_image_token/replace_video_token/create_mm_token_type_ids in
both cache (geometry-only, no media decode) and runtime paths, so the
two paths run the identical HF code rather than two independent
reimplementations. Mixed image+video in one sample raises
NotImplementedError instead of silently dropping video.
- xtuner/v1/datasets/data_item.py: Glm53VLDataItem.
- xtuner/v1/data_proto/sequence_context.py: pixel_values_videos /
video_grid_thw / mm_token_type_ids fields, wired through
__init__/split/cat/to/copy/data.
- xtuner/v1/datasets/collator.py: build_text_ctx_labels now also
returns kept_indices (which original instances survived pack-level
truncation); glm53_vl_sft_collator built on top.
Root-caused two real discrepancies via direct package inspection /
minimal repro rather than trusting the design doc's claims: (1)
transformers 5.17.0 has no Glm5NextVideoProcessor.get_number_of_video_
patches despite the design doc and an independent peer check both
asserting it exists -- verified missing via hasattr/grep across the
installed package, and confirmed the intended cache-path geometry
prediction (Glm5NextVideoProcessor.sample_frames + the module-level
smart_resize -- the same two public calls the real preprocessing path
makes) reproduces the real video_grid_thw bit-for-bit; (2) an earlier
version expanded the whole <|begin_of_image|>...<|end_of_image|> marker
instead of just the inner <|image|> token, undercounting num_tokens by
2 per media item -- caught by an end-to-end diff against a real
processor() call, fixed by only replacing the bare inner token.
doc/xtuner_glm5p3flash_design.md / .py: synced the F1.c reference
table and the cache-path docstring to reflect finding (1) above --
these still claimed both get_number_of_image_patches and
get_number_of_video_patches were exposed pure-geometry APIs, which was
true only for the image side; the video side needs the sample_frames +
smart_resize fallback described above. This was an oversight in the
original commit: the code and this commit message already reflected
the real finding, but the design doc's own reference tables didn't get
updated to match at the time.
Known gap: DataloaderConfig.build_collator() doesn't yet register
glm53_vl_sft_collator (needs image_token_id/merge_unit threaded from
the tokenize_fn instance) -- left for F6's end-to-end wiring.
Test plan: tests/datasets/test_glm53_chat.py (6, byte-exact vs HF
jinja), test_glm53_vl_collator.py (9, includes build_text_ctx_labels
regression coverage for qwen3_vl_sft_collator/intern_s1_vl_sft_collator
since they share the builder), test_glm53_vl_tokenize_fn.py (10,
includes two tests asserting input_ids/mm_token_type_ids match a real
end-to-end processor() call token-for-token). All 25 pass, CPU only.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
`glm53_vl_sft_collator` landed with F1 but was never reachable from a training config: `DataloaderConfig.build_collator()` only knew the three older collators by name, and the dotted-path fallback could not supply the two extra arguments this collator needs (`image_token_id`/`merge_unit`, used to re-check the placeholder<->patch-count invariant after pack-level truncation). Any GLM-5.3-Flash VL config therefore died in the dataloader with a `TypeError` before the first step -- the F1 "known gap" that blocked end-to-end VL training. Registers the name, and adds a generic `collator_kwargs` that is bound into the collator partial for any collator needing more than the four arguments the dataloader supplies itself. The name->function dispatch collapses into a lookup table so every branch, including the dotted-path one, goes through the same binding. Test Plan: `test_glm53_collator_is_reachable_by_name_from_dataloader_config` builds the collator through `DataloaderConfig` and calls it with only the dataloader-supplied arguments; the other four names plus a dotted path still resolve. tests/datasets/test_glm53_vl_collator.py 7/7. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`xtuner.v1.datasets.mllm_tokenize_fn` re-exported it but the package root did not, unlike every other VL tokenize-fn config, so a training config had to reach into the submodule by path. Makes the GLM-5.3-Flash VL entry point look like Qwen3-VL's and InternS1's. Test Plan: `from xtuner.v1.datasets import Glm53VLTokenizeFnConfig, Glm53VLTokenizeFunction`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The class documented that GLM-5.3-Flash renders through `render_glm53_chat` instead of a `CHAT_TEMPLATE_MAP` entry, and `build()` never reads `chat_template` -- but the field is declared required on `BaseMLLMTokenizeFnConfig`, so constructing the config the documented way raised `ValidationError: chat_template Field required` before any data was touched. Defaults it to the template name the renderer implements, which keeps the field's declared type while making the documented construction work. Test Plan: `Glm53VLTokenizeFnConfig(processor_path=..., max_length=...)` constructs; exercised end to end by the VL SFT config in the F6 layer. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The VL e2e matrix showed `tgs` at roughly half of `seqlen_tgs` on every profile, where the text matrix has them within 0.05%. The two differ only by padding (`mask.sum()` vs `mask.numel()`), so VL packs were ~50% padding -- one step packed 131 real tokens into a 16384 slot. The dataset was large enough to fill them (2.38M tokens against 2.62M slots), so this was not starvation. Comparing each sample's cached `num_tokens` (which the packer budgets with) against its runtime length found 17% of samples cached at 4096 but training as **33-token fakes**. Two defects, both in the cache/runtime split, compounded: 1. **Frame-folder videos were indexed against the original video.** Both paths built `VideoMetadata` from `origin_video_length`/`origin_fps` (182 frames @30fps), but a frame folder holds the video *after* resampling -- `processed_video_length` frames at `processed_fps` (23 @4fps). Runtime indexed a 23-file folder with original-video frame numbers and raised `IndexError: list index out of range`; the dataset swallowed it and substituted a fake sample. Sampling now uses the processed sequence for frame folders, qwen3_vl's `calc_frame_info` rule 3, which also makes the frame timestamps correct. The folder/file decision is read from the path string (the trailing `/` convention qwen3_vl's reader keys on), not the filesystem, so the cache path -- which never touches the media -- reaches the same answer as the runtime path. 2. **Over-length visual samples were packed at their truncated length.** With the index fixed, the same samples expand to 14,352 placeholders, far past `max_length=4096`. Truncating text is harmless, but truncating *through* a visual span leaves fewer placeholders than patches; the cache path reported the truncated 4096 anyway, so the packer budgeted for it and runtime then rejected it -- again becoming a fake. The cache path now reports such samples as damaged (`num_tokens=0`) so they are filtered before packing, the same check qwen3_vl's cache path makes ("如果图片被截断，则该数据丢弃"). Both escaped the existing tests for concrete reasons: `_video_item` only set `origin_*` on a folder holding exactly that many frames, so origin and processed coincided; and the truncation test only asserted the *runtime* path raises -- which is exactly the behavior that gets silently turned into a fake sample, while the cache path, where the damage is decided, was never checked. Test Plan: - `test_glm53_frame_folder_samples_the_processed_frames`: a folder resampled from a longer original (180 frames @30fps -> 6 @2fps); cache and runtime agree and runtime does not raise. - `test_glm53_visual_truncation_is_dropped` now also asserts the cache path returns `num_tokens == 0`, keeping the runtime rejection as a last line of defence. - Both fail on the previous code with the exact production errors (`IndexError: list index out of range`; `assert 20 == 0`) and pass with the fix. tests/datasets/test_glm53_vl_tokenize_fn.py 11/11. - Over the whole ci_vl corpus: 0 cached/runtime length mismatches across the 112 kept samples, with 24 over-length samples filtered before packing (previously 10 of the first 60 mismatched). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A pack that holds both an image sample and a video sample used to enter the tower once per modality. The tower's FSDP mesh is the world group and the language tower's is the dp group, so ranks disagreed on the collective sequence and baseline hung on step 4. One concatenated call keeps every rank in lockstep, and each grid row stays its own attention segment.
Three precision nets over the composed stack: an AdamW step through the public training path; seven sampled gradients (first/last KDA + mHC + final norm) on the five-layer crop under 2-rank FSDP, within 6% relative error of HF; and a 2-rank multimodal FSDP backward where one rank is text-only while the other carries both image and video, proving ranks join the same vision/projector collectives.
heqs1013
reviewed
Oct 12, 2026
…bedding caches (InternLM#2148) * [Fix] Allow multiple logical depths for the shared GLM5.3 MTP layer The checkpoint contains one physical MTP layer, but MTPBlock can reuse it at multiple logical prediction depths. Restricting num_layers to one rejects valid seven-depth configurations. Require shared weights and positive logical depth while still constructing one physical layer and retaining the layers.45 checkpoint mapping. (cherry picked from commit df9924d) * [Fix] Preserve FP32 MTP TV loss on empty supervision shards Non-empty TV losses are FP32, while an unsupervised BF16 shard previously created a BF16 zero before the distributed SUM reduction. Compute empty-branch sums in FP32 so every rank uses the same loss dtype. Retain zero-gradient edges to the LM head and each draft hidden state. (cherry picked from commit 2025835) * [Test] Cover shared GLM5.3 MTP depths and empty TV shards Adapt the MTP coverage from 87bfd92 to the F6 model: build CE loss contexts and run forward/backward at one and seven logical depths. Check the single physical layer, checkpoint mapping, and accumulation of gradients across recurrent uses of its shared parameters. Cover FP32 losses and zero-gradient edges on empty BF16 supervision shards, with both attached and detached LM heads. Add a two-rank Gloo regression for seven-depth CE and TV with one unsupervised SP rank. Validation: 12 CPU regressions passed. The two full-model CUDA cases were collected but skipped because CUDA was unavailable. * [Feat] Cache local packed MTP future embeddings Add an opt-in local future embedding path that gathers once per microbatch and materializes independent local buffers for each prediction depth. Preserve pack boundaries, checkpoint input lifetimes, supplied-embedding gather semantics, and the token-ID fallback. Port only optimization A from c25acadb onto the GLM5.3 MTP PR base; exclude shared CE head gradients, TV precomputation, teacher reuse, profiling and private training recipes. Default use_local_future_embeddings remains False. Validation: 67 single-process CPU checks and two loopback Gloo regressions pass for local buffer values/storage, checkpoint gradients, shared GLM depths, mixed empty supervision and SP boundaries. Ruff and formatting checks pass. * [Feat] Enable local MTP future embeddings by default Validate CUDA packed cache values and gradients, two-rank NCCL seven-depth TV checkpoint backward, and real GLM53 CE/TV supplied-embedding losses and gradients. Keep explicit legacy opt-out and token-ID fallback.
JT-Ushio
changed the title
JT-Ushio
changed the base branch from
feat/glm53flash-f3-kda
to
feat/glm53flash-f6-text-moe
October 12, 2026 13:32
jayhenry
reviewed
Oct 12, 2026
jayhenry
force-pushed
the
feat/glm53flash-f6-text-moe
branch
2 times, most recently
from
October 12, 2026 19:22
e006658 to
48097aa
Compare
jayhenry
reviewed
Oct 13, 2026
jayhenry
merged commit Oct 13, 2026
31d0539
into
InternLM:feat/glm53flash-f6-text-moe
1 of 3 checks passed
jayhenry
pushed a commit
that referenced
this pull request
Oct 13, 2026
jayhenry
pushed a commit
that referenced
this pull request
Oct 14, 2026
jayhenry
pushed a commit
that referenced
this pull request
Oct 14, 2026
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.
Fix KDA kernel switching during eval-mode checkpoint replay and unsupported Tensor options that were ignored or raised ambiguous boolean errors. Preserve FP32 convolution/gate gradients and retain Ulysses as the default, with FLA sequence parallelism available via
sp_impl="fla".Base:
feat/glm53flash-f6-text-moeat48097aa4b2102f052e309c7c0945a14894aa0d51(#2111). F6 includes F6.1 (#2148); this PR is F6.2 and targets F6 directly.Changes
eval()+torch.inference_mode();eval()+no_grad()uses chunked KDA.output_final_state=Falseseparately.conv_kernel_size - 1local tokens; Ulysses retains its fixed-layout gate prefix.XTUNER_HF_IMPLreference entrypoints. Pin convolution masters through the F6 FP32 parameter configuration. This also exports convolution tensors as FP32; BF16 inference loads restore the forward-rounded values.Validation
Validated commit
446376con 4 H200 GPUs throughPJT-via-Fudan, using PyTorch 2.9.1+cu128 and FLA 0.4.2: 141 passed, 4 deselected (SP8 requires eight GPUs), in 226.42 seconds. Job:pr2153-integrated-446376c-94730520. Pydantic checks, local lint, and Linux-targeted mypy passed. The follow-up commita50a578changes docstring formatting only; executable ASTs are unchanged. GitHub lint passed ona50a578. Full exported-model inference and training trajectory parity were not tested.The Read the Docs checks fail on
446376cand on the F6 base; build logs are inaccessible in this session.Numerical controls: standalone XTuner chunk gradients match native FLA exactly on the previously failing FP32 safe-gate cases; both show approximately 0.5% gate-gradient error against the recurrence. The standalone safe-gate FP32
dgoracle budget is 1%, with a separate 0.01% native-parity budget. Direct sequence-SP tests retain 0.3% native parity for FP32dq/dk/dv, use 0.5% for FP32dg/dbeta, and check both implementations against a 0.5% recurrence budget. BF16 gradient budgets remain 2%; isolated zero gradients must match exactly.Existing benchmark in
c9f5ebf(not rerun here): one KDA layer, 8 H200s, SP8, 128K global tokens, hidden size 4096, 64 heads of dimension 128, BF16 activations and FP32 convolution/gate masters. With 3 warmups and 10 timed iterations, median maximum-rank forward/backward time was 76.00 ms for Ulysses and 44.35 ms for FLA (1.71×). These KDA-only timings exclude cold compilation, outer checkpointing, the remaining model, loss/optimizer, and outer parameter-gradient synchronization; they do not establish full-model throughput or trajectory parity.Provenance and integration
d3e2d300c98c63974c39ff3372e083d247bf15e220586b43f014ab00The four extracted commits retain their
Original-Committrailers from the training-correctness branch. The F6fp32_keys_patternpreserves convolution masters and saves them as FP32, while convolution forward still uses the activation dtype. Functional changes are limited to KDA. Two inherited sparse-MLA files receive docstring formatting to pass repository-wide lint.