[GLM 5.3 Tests] Reproduce trainer batch-size logprob sensitivity

[hershg]

What this adds

• A trainer-only control: score one fixed sequence, two identical copies together, then the original again. Save raw token scores and batch/repeat/duplicate differences; no invented pass/fail tolerance.
• Reusable runner and tests in examples/model_checks; pinned Qwen fixtures there and GLM-5.3 fixtures in examples/tinker/glm5p3. No optimizer, inference, or weight publication—this does not test router-state drift or LoRA delivery.
• Explicit preflight/zero-adapter checks and Ray cleanup even when report storage fails.

Evidence

Runtime	Mean / max: one row vs duplicated batch	Duplicate-row diff	Repeat diff
Qwen TP1, historical	5.34e-8 / —	0	0
Qwen TP8, historical	0.030262 / —	0	0
GLM TP8/EP8, historical	0.217588 / 3.24219	0	0
GLM TP8/EP8, current combined stack	0.217588 / 3.24219	0	0

The current GLM rerun completed on one B300 with immutable image sha256:4b7d4ae3e41561bd691862bfed20ee87eb135a8e14b0bee6f6fd6d836d88d26a and integrated SkyRL 09d6f7b83d965b6dc79e5524b920b7a83000a9f8. It scored 64 positions: initialization 497.27s, single row 92.53s, duplicated batch 2.38s, repeat 2.08s. The current PR cleanup changes layout, errors, and cleanup—not the fixed fixture, scoring path, or comparison semantics.

This is an unresolved characterization, not a calibrated learning failure: deterministic repeat and duplicate-row agreement rule out random execution noise, while the large one-row/two-row change remains specific to trainer batch shape. No universal acceptance threshold is asserted here.

Validation / reproduction

Current cleanup: 13 pure tests passed, Ruff/Black/gitleaks passed; renamed fixtures are byte-identical. The expanded 24-test native suite cannot collect on the devbox because its local CUDA environment lacks cuDNN; the live B300 rerun above exercises the actual model path.

uv run --isolated --no-project --with pytest python -m pytest --noconftest \
  examples/model_checks/tests/test_batch_sensitivity.py -q

Inside the pinned prebuilt CUDA image:

uv run --active --no-sync --extra tinker --extra megatron \
  python examples/model_checks/run_batch_sensitivity.py \
  --model /models/304b8051cfb2b260b61ce0cbe330e02a98e73639 \
  --backend-config examples/tinker/glm5p3/batch_sensitivity_config.json \
  --fixture examples/tinker/glm5p3/batch_sensitivity_tokens.json \
  --output-dir /artifacts/glm-batch-sensitivity

[gemini-code-assist]

Code Review

This pull request introduces a batch sensitivity diagnostic tool to compare token scores across different batch sizes on a Megatron trainer, including configuration fixtures and tests for Qwen3-8B and GLM-5.3 models. The review feedback focuses on improving robustness and error handling by replacing bare assert statements with explicit exceptions, adding defensive checks for potential None values in tokenizer padding, and handling missing keys in dictionary lookups.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge, with no concrete changed-code defect identified.

The configured duplicate rows remain together in one token-based microbatch, worker output preserves per-sample scores, and the added validations consistently enforce the diagnostic’s stated input and scoring contracts.

Important Files Changed

Filename	Overview
examples/model_checks/run_batch_sensitivity.py	Implements the isolated Megatron scoring workflow, validates topology and inputs, records all token scores, and writes incremental reports.
examples/model_checks/batch_sensitivity.py	Provides finite, aligned token-score comparisons and separates the three diagnostic error categories.
examples/model_checks/tests/test_batch_sensitivity_runtime.py	Verifies batch equality, adapter validation, production scoring dispatch, and output-shape rejection.
examples/model_checks/tests/test_batch_sensitivity.py	Covers comparison calculations and malformed or nonfinite score handling.
examples/model_checks/fixtures/qwen3_8b_config.json	Pins the single-GPU Qwen Megatron topology and batching settings used by the diagnostic.
examples/tinker/glm53/batch_sensitivity_config.json	Pins the eight-GPU GLM tensor/expert-parallel topology and model-specific transformer settings.

Sequence Diagram

sequenceDiagram
    participant Runner
    participant Ray as Ray Actor Group
    participant Worker as Megatron Worker
    participant Report
    Runner->>Ray: Score one fixed sequence
    Ray->>Worker: "forward(loss_fn=None)"
    Worker-->>Runner: Single-row token logprobs
    Runner->>Ray: Score two identical copies together
    Ray->>Worker: "forward(loss_fn=None)"
    Worker-->>Runner: Two token-logprob rows
    Runner->>Ray: Score original sequence again
    Ray->>Worker: "forward(loss_fn=None)"
    Worker-->>Runner: Repeat token logprobs
    Runner->>Report: Compare batch-size, duplicate, and repeat errors

Loading

_{Reviews (1): Last reviewed commit: "test: isolate trainer batch-size logprob..." | Re-trigger Greptile}