[GLM 5.3 Tests] Check trainer/inference LoRA logprob agreement

[hershg]

LoRA publication check

This adds one focused invariant: a LoRA update produced by the trainer must reach inference without silently loading stale or numerically divergent weights.

The check scores identical fixed tokens at three points:

Stage	Comparison	Required result
A	Zero-init trainer vs inference	mean < 0.05; max < 0.5
B	Updated trainer vs still-stale inference	mean >= 0.05 for the deterministic test stimulus
C	Updated trainer vs inference after publication	mean < 0.05; max < 0.5

Mean and maximum error are independent gates. The localized-corruption regression demonstrates why: one token with error 9.0 has mean error 0.046875 across 192 positions, so it passes the mean budget and must fail the maximum budget.

Evidence

Model / path	A: zero mean / max	B: stale mean	C: published mean / max	Result
Qwen3-0.6B, native runner	0.016145 / —	0.100314	0.020233 / —	Pass on the original mean contract
Qwen3-8B, native runner	0.023879 / —	0.200157	0.021330 / —	Pass on the original mean contract
GLM-5.3, two-node R3 target-execution receipt	0.005877 / 0.424854	0.440516	0.030078 / 0.458789	Pass

The GLM result uses zai-org/GLM-5.3-BF16, a TP8 vLLM inference node and a TP8/EP8 Megatron training node, rank-32 LoRA over all configured target families, identical tokens and scored positions, and replayed R3 expert routes. All eight target families were present and exercised. This supersedes the earlier pre-R3 failure: without replaying identical sparse-expert routes, trainer/inference comparison was not a valid same-computation control.

A separate current-main cold lifecycle replay also passed after one real GSPO optimizer update: zero mean/max 0.006000 / 0.311349, unchanged pre-publication sampler, published mean/max 0.000706 / 0.015649, held-out mean 0.000761, finite loss and gradients, matching transfer digests, checkpoint, unload, trainer offload, and release. A normal optimizer update can be smaller than the 0.05 synthetic stale-detection threshold; this receipt therefore checks unchanged-before-sync plus agreement-after-sync rather than falsely requiring a large learning step.

The public Tinker SDK path was independently exercised on Qwen3.5-4B through ServiceClient -> create_lora_training_client -> save_weights_and_get_sampling_client: XID 1052900, zero mean 0.01869, stale mean 1.099, republished mean 0.00562, all 192 scores finite.

Scope

This PR supplies reusable comparisons plus native and public-SDK runners. CPU coverage includes alignment, nonfinite scores, stale/wrong weights, insufficient stimulus, localized tail corruption, deterministic perturbation, and cleanup. It does not claim batching invariance, optimizer stress, learning quality, or universal bit-exact backend agreement.

The current GLM runtime proof also depends on the separately reviewed runtime fixes for remote disaggregated publication, shared-expert TP-overlap scaling, and MLA kv_b_proj LoRA execution; those changes are intentionally not folded into this testing PR.

[hershg]

Moved active review to hershg#2, which genuinely targets PR1’s branch and shows only the child diff. This closes the misleading combined-diff view, not the work. GitHub cannot target a contributor-fork branch from an upstream PR. After #2181 merges, rebase the child onto upstream main and reopen this PR for upstream landing. No code was merged; native GPU validation remains pending.

[gemini-code-assist]

Code Review

This pull request introduces numerical checks and adapter-only perturbation utilities for GPU LoRA testing, along with standalone Megatron LoRA diagnostics and CPU tests. The reviewer's feedback provides valuable robustness improvements, including wrapping JSON serialization to prevent masking assertion errors, ensuring device alignment during log probability comparisons, generating random noise in float32 to avoid PyTorch compatibility issues on CPU, and casting response mask sums to long integers to prevent float-to-integer comparison issues.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge based on the reviewed changes, with no concrete blocking or independently actionable non-blocking defect identified.

Diagram

sequenceDiagram
    participant R as Diagnostic runner
    participant T as Megatron trainer
    participant S as Inference sampler
    R->>S: Score base model
    R->>T: Publish zero-initialized LoRA
    R->>S: Score zero adapter and repeat
    R->>T: Score zero adapter and repeat
    R->>R: Check initial agreement and noise
    R->>T: Perturb trainable LoRA parameters
    R->>T: Score updated trainer
    R->>S: Score unpublished adapter
    R->>R: Verify sampler remains unchanged
    R->>T: Broadcast updated LoRA
    T->>S: Load skyrl-lora weights
    R->>S: Score published adapter
    R->>R: Check parity, update magnitude, and delta agreement

Loading

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

^{Reviewed by Cursor Bugbot for commit f9b1b1c. Configure here.}

[j316chuck]

Two remaining changes before this is a correctness check:

1. Add a localized-corruption regression and gate tail error. With 192 positions, set zero/stale=-2.0, trainer_updated=-1.9 everywhere, and updated=trainer_updated except one position at -10.9. Today check_updated_adapter(report, 0.05) passes because mean_abs=0.046875, despite max_abs=9.0. Add this as a failing regression, then require an explicit max_abs budget (for initial and updated parity) so one bad token cannot hide in the mean.

2. Run publication through the public Tinker sync path. Make the live runner use tinker.ServiceClient -> create_lora_training_client -> save_weights_and_get_sampling_client, apply a fixed forward_backward + optim_step, score the existing sampler while stale, then call save_weights_and_get_sampling_client again and score the published model. Reuse skyrl/tinker/sdk_logprobs.py. The direct broadcast_to_inference_engines call bypasses WorkerDispatch.save_weights_for_sampler behavior such as adapter/model-id selection and weight-version/cache handling.

Already resolved on current head:

3. The isolated pytest command now installs pytest with --with pytest.
4. Native trainer and sampler paths reject NaN/Inf before inserting scores into the JSON receipt, with regression coverage for both.

[j316chuck]

Tentative approval based on a live Qwen/Qwen3.5-4B Megatron run against exact head 1e0eddb (TP4 trainer + 2xTP2 vLLM) through the public Tinker SDK path.

With one lr=1e-3 optimizer update, all 192 probe scores were finite: initial trainer/sampler mean_abs=0.01869, repeat and withheld-publication drift=0, stale mean_abs=1.099, and re-published mean_abs=0.00562. Image digest: sha256:fe8bc62b34b2cb28a5a42a39fa07cb6085e5fdd22dc658b2eb0004193f121de5; deployment XID: 1052900.

Caveat: a deliberately large 2x lr=1e-2 update exceeded the 0.05 absolute-parity budget (mean_abs=0.185, max_abs=1.172) despite cosine=0.99991. Please retain the max/tail-error follow-up; this approval is scoped to the realistic passing update path.