Skip to content

[model] fix: handle separate KV normalization in GLM TileLang attention - #37

Open
yushengsu-thu wants to merge 3 commits into
bridgefrom
codex/lora-effective-weight
Open

[model] fix: handle separate KV normalization in GLM TileLang attention#37
yushengsu-thu wants to merge 3 commits into
bridgefrom
codex/lora-effective-weight

Conversation

@yushengsu-thu

@yushengsu-thu yushengsu-thu commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Newer GLM DSA specs normalize the KV latent in a separate module. The TileLang absorb path previously assumed a fused projection layernorm, accessing an absent attribute and attempting to normalize the latent again. Read the fused layernorm weight only when present, applying KV RMSNorm exactly once while retaining the existing differentiable LoRA weight folding.

Validation on 3dabfd2ae840cbb7e9080d3b4e398c0e5eb7843b:

  • Full 78-layer GLM-5.3 LoRA DAPO, 8192-token context, TP8/EP16/DP2 on 2 nodes / 16 H200s: one forward/backward/optimizer/save step completed. Gradient norm 0.211709; all 59,073 exported adapter tensors were finite, including 156 nonzero KV-up A/B tensors across all 78 layers. Measured run.
  • FP8 rollout versus BF16 trainer: token logprob MAE 0.027405, P95 0.153260, maximum 5.331019. The sample-mean gap 0.002939 passed the original 0.03 CI gate. Tokens, rollout logprobs and raw/normalized rewards matched the saved rollout through the DP handoff.
  • The run later exhausted host RAM during trainer offload, before post-training publication. A new run uses one cross-node rollout engine to reduce duplicate CPU backups; consecutive cycles remain under validation.
  • Changed-file pre-commit passed. Whole-repository lint has unrelated failures; strict mypy was unavailable locally. Full Bridge CI was not run.

Companions: Miles #3174, SGLang #38760.

The v10 cross-node retry used SGLang f6f03206a74bce5910d04f3581a423a68beb25d5, which fixes a reproduced peer-local port precheck false positive. Its11 PortArgs tests and changed-file pre-commit pass; full GPU evidence above remains scoped to the preceding SGLang head.

V10 was externally stopped through the Ray stop API at2026-09-10 06:12:46UTC, after selecting12/16 samples and before its first optimizer step. The reward guard did not trigger and no preceding training failure was observed; the stop caller is unidentified. No job is currently running. Host-memory peaks during startup/sampling were873.23GiB and870.13GiB, but the post-training offload peak and consecutive publication cycles remain unverified.

Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
@yushengsu-thu yushengsu-thu changed the title [peft, model] fix: support LoRA weights and split KV norm in GLM DSA [model] fix: handle separate KV normalization in GLM TileLang attention Sep 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant