[Bug] GLM-5.3 hybrid absorbed MLA silently bypasses KV-up LoRA gradients via direct .weight access

Start with the native _get_kv_up_weights in megatron/core/transformer/experimental_attention_variant/absorbed_mla.py, the GLM factory in src/mcore_bridge/model/mm_gpts/glm5_next.py, and the LoRA wrapper in src/mcore_bridge/tuners/lora.py. Run reproduce_kv_up.py to confirm that absorption leaves the adapter B gradient and update absent while the base weight stays frozen. Done means the factory-selected native class receives a differentiable effective KV-up weight, with regression coverage for K/V updates, unchanged frozen weights, precision, tensor parallel, recompute, and export/reload behavior.

Checklist / 检查清单

• I searched the repository's public issues and PRs for absorbed MLA / KV-up / LoRA gradient reports. I did not find an exact duplicate in that search; related prior work is linked below.

Bug Description / Bug 描述

When GLM-5.3-Flash uses the native Megatron AbsorbedMLASelfAttention through the hybrid layer factory, KV-up absorption reads linear_kv_up_proj.weight directly. With the current adapter wrapper contract, that property exposes the frozen base weight, while the trainable LoRA contribution is applied in the wrapper's forward().

The absorbed K/V computations therefore have no autograd connection to the KV-up adapter A/B parameters. Training can still run, loss can change, and checkpoints can be saved because other parameters/activations remain in the graph. This is a training correctness report, not an inference-cache or performance report.

Expected: the effective K and V up-projection weights used by absorption include the active adapter contribution differentiably, without modifying the frozen base.

Observed: in the CPU isolation below, ordinary LoRA forward updates B; the upstream absorption getter bypasses forward, leaves A/B gradients absent, and B stays zero after two optimizer steps. A differentiable effective-weight reference updates B.

A downstream GLM hybrid smoke run also had all 11 KV-up B tensors remain zero, in both the native distributed checkpoint and the exported adapter. A bridge-subclass-only workaround did not cover the actual native class selected by this model. Moving the adapter-aware behavior onto the actual path, with associated precision/TP handling, resolved our downstream check. This is not a claim that every MLA/LoRA implementation is affected.

How to Reproduce / 如何复现

Original environment: Python 3.12, PyTorch 2.11.0+cu130, PEFT 0.20.0, Transformer Engine 2.18.0, BF16 frozen base / FP32 adapters, TP4/EP4/CP1, sequence parallel and full recompute. The downstream setup also contains custom NoRA support; the CPU reproduction below deliberately uses ordinary stock PEFT LoRA and no NoRA patch.

Pinned sources:

• mcore-bridge: 63b3570fa18bbdbaeaa3d07a57bf50092815c216
• ms-swift: 86989ac34b67765fc4f8a48c87d3ad9d22f53540
• Megatron-LM: c4220af3a62e851585afdac6f5aa7496163afdaf

Relevant code:

• GLM hybrid factory: copies hybrid_stack_spec and configures its DSA submodules. In our factory probe, spec.submodules.dsa_layer.submodules.self_attention.module is the native Megatron class, not mcore_bridge.model.modules.absorbed_mla.AbsorbedMLASelfAttention.
• Native KV-up getter: accesses .weight and reshapes/splits it.
• Adapter forward: adds the active adapter contribution in forward().

Small CPU isolation, no model/data/GPU required: install compatible PyTorch and peft==0.20.0; save the script below as reproduce_kv_up.py, then run:

curl -fL https://raw.githubusercontent.com/NVIDIA/Megatron-LM/c4220af3a62e851585afdac6f5aa7496163afdaf/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py -o absorbed_mla.py
python reproduce_kv_up.py --absorbed-source absorbed_mla.py

The script extracts and executes the actual upstream getter unchanged, on a tiny standard PEFT layer. It does not copy a replacement getter, load the model, import the full Megatron runtime, or claim to reproduce complete GLM training. Its differentiable reference is a control, not a production patch. A separate factory probe and downstream checkpoint inspection provide the model-path evidence.

Observed results (two optimizer steps, B initialized to zero):

The frozen base is unchanged in all three cases. The reference's zero hook count illustrates that bypassing forward is not itself wrong; the issue is failing to include the adapter contribution in the effective weight.

"""CPU-only isolation of upstream absorbed MLA's KV getter with standard PEFT LoRA.

No model, CUDA, private data, NoRA patch, or full Megatron import is required.
The getter is extracted unchanged from an explicitly supplied upstream source file.
This tests the getter/adapter contract, not a full GLM training forward.
"""
import argparse
import ast
import copy
import hashlib
import json
from pathlib import Path
from types import SimpleNamespace

import torch
import peft
from peft import LoraConfig, get_peft_model


def load_getter(path):
    source = path.read_text()
    tree = ast.parse(source)
    cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == 'AbsorbedMLASelfAttention')
    node = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == '_get_kv_up_weights')
    namespace = {'torch': torch}
    exec(compile(ast.Module(body=[node], type_ignores=[]), path.name, 'exec'), namespace)
    return namespace[node.name]


class ToyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.kv_up = torch.nn.Linear(3, 8, bias=False)

    def forward(self, x):
        return self.kv_up(x)


def wrap():
    model = get_peft_model(ToyModel(), LoraConfig(r=2, lora_alpha=2, lora_dropout=0., target_modules=['kv_up']))
    return model.base_model.model.kv_up


def holder(layer):
    return SimpleNamespace(_uses_combined_kv_up_projection=True,
                           num_attention_heads_per_partition=2,
                           config=SimpleNamespace(kv_lora_rank=3, qk_head_dim=2, v_head_dim=2),
                           linear_kv_up_proj=layer)


def absorbed_loss(k, v):
    # Trainable upstream activations allow backward to succeed even when the
    # frozen KV weight has no graph connection to A/B: the silent failure case.
    q = torch.randn(5, 1, 2, 2, requires_grad=True)
    latent = torch.randn(5, 1, 2, 3, requires_grad=True)
    return (torch.einsum('sbhd,hdk->sbhk', q, k).square().mean()
            + torch.einsum('sbhk,hdk->sbhd', latent, v).square().mean())


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--absorbed-source', type=Path, required=True)
    parser.add_argument('--expect', choices=['bypass', 'connected'], default='bypass')
    args = parser.parse_args()
    torch.manual_seed(42)
    getter = load_getter(args.absorbed_source)
    seed_layer = wrap()
    cases = {}
    for mode in ('linear_forward', 'upstream_absorption', 'differentiable_reference'):
        layer = copy.deepcopy(seed_layer)
        a, b = layer.lora_A['default'].weight, layer.lora_B['default'].weight
        base_before = layer.get_base_layer().weight.detach().clone()
        optimizer = torch.optim.SGD([a, b], lr=0.1)
        hook_calls = []
        hook = layer.register_forward_hook(lambda *unused: hook_calls.append(True))
        grad = None
        for _ in range(2):
            optimizer.zero_grad(set_to_none=True)
            if mode == 'linear_forward':
                loss = layer(torch.randn(5, 3, requires_grad=True)).square().mean()
            elif mode == 'upstream_absorption':
                k, v = getter(holder(layer))
                loss = absorbed_loss(k, v)
            else:
                # Standard LoRA reference, deliberately not a proposed production fix.
                w = layer.get_base_layer().weight + (b @ a) * layer.scaling['default']
                w = w.view(2, 4, 3)
                loss = absorbed_loss(w[:, :2], w[:, 2:])
            loss.backward()
            grad = b.grad
            optimizer.step()
        hook.remove()
        cases[mode] = {
            'forward_hook_calls': len(hook_calls),
            'B_grad_is_none': grad is None,
            'B_grad_nonzero': 0 if grad is None else int(grad.count_nonzero()),
            'B_nonzero_after_two_steps': int(b.count_nonzero()),
            'A_grad_is_none': a.grad is None,
            'base_unchanged': torch.equal(base_before, layer.get_base_layer().weight),
        }
    assert cases['linear_forward']['B_nonzero_after_two_steps'] > 0
    assert cases['differentiable_reference']['B_nonzero_after_two_steps'] > 0
    assert all(c['base_unchanged'] for c in cases.values())
    bypass = cases['upstream_absorption']['B_nonzero_after_two_steps'] == 0
    assert bypass == (args.expect == 'bypass'), cases
    print(json.dumps({'torch': torch.__version__, 'peft': peft.__version__,
                      'source_sha256': hashlib.sha256(args.absorbed_source.read_bytes()).hexdigest(),
                      'getter_only_not_full_model': True, 'cases': cases,
                      'result': 'bypass_reproduced' if bypass else 'adapter_connected'}, indent=2))


if __name__ == '__main__':
    main()

Additional Information / 补充信息

Current upstream check, 2026-11-14: I repeated the getter-only CPU test against Megatron-LM 4b11b8168e9e518c0e00bc691d0265cccf1820e6, with the same result. At mcore-bridge 8f16307c3ea77ab44a5a2800a5ce6abcb9a4d71c, the inspected GLM factory file and LoRA wrapper file match those at our pinned bridge revision; the wrapper has no adapter-aware weight override or get_weight_for_absorption method. This is source inspection plus a getter-level reproduction, not a full GLM rerun of the latest upstream combination.

Factory-only probe result:

{
  "selected_class": "megatron.core.transformer.experimental_attention_variant.absorbed_mla.AbsorbedMLASelfAttention",
  "is_native_attention": true,
  "is_bridge_subclass": false,
  "model_weights_loaded": false
}

Downstream fix evidence: after repairing the actual native path and preserving differentiable FP32 adapter contributions, our LoRA r32/alpha64 and NoRA r128/alpha128 runs each reached 46/46 steps and all 11 MLA layers passed separate K/V nonzero-update checks. These are existing downstream results, not clean-upstream validation or proof of model quality. Other local patches were present; this report does not claim the complete downstream patch stack is ready to merge upstream.

Suggested ownership/design question: should mcore-bridge expose an adapter-aware effective-weight contract for absorption, ensure the hybrid factory routes through the relevant bridge implementation, or coordinate a generic accessor with Megatron Core? We would prefer a maintainer-approved contract rather than adding a PEFT-specific dependency to Core.

A fix should be tested on the class returned by the actual GLM factory, not only on the bridge subclass. Regression coverage should include:

• K and V B gradients/updates, and A gradients after B becomes nonzero; zero A gradient on the first step of standard zero-B initialization is expected.
• Frozen base unchanged; disabled/merged/unadapted paths consistent.
• TP2/TP4 and sequence parallel on/off, with replicated-factor gradients reduced exactly once.
• BF16 activations with FP32 adapters: avoid losing small updates through premature materialization of the effective weight in BF16; check both K absorption and V projection.
• Recompute and export/reload consistency; explicit handling or rejection of unsupported dropout/bias cases.

Related prior work: NVIDIA-NeMo/Megatron-Bridge #4198 adds LoRA-aware weight access for MLA absorption in a different adapter stack. This is relevant precedent; I am not claiming the underlying absorption/LoRA interaction is newly discovered. mcore-bridge #144 introduced absorbed MLA support, while #216 concerns a different selective-recompute argument-forwarding issue.

I can provide the factory probe, sanitized per-layer update checks, and downstream reference changes if useful. No private training samples or model weights are needed for the small reproduction above.