LongStraw Training Architecture

This checkout is the GLM-5.2 instantiation of LongStraw's broader objective-aware, architecture-aware resident-state virtualization, response-replay, and distributed-gradient execution contract. The resident manager is only one part: the objective selects state lifetime, replay controls recomputation, and gradient finalization controls whether a distributed update is semantically closed.

Design objective

Conventional GRPO recomputes the full prompt for every completion and retains a full-sequence autograd graph. At multi-million-token context lengths, that working set is the limiting factor. This implementation separates the read-only prompt state from the short trainable response path:

                         once per prompt
tokens  ->  CP prefix forward  ->  compact resident state
                                      |       |       |
                                      v       v       v
                                  response  response  response
                                    replay    replay    replay
                                      \       |       /
                                       accumulated LoRA grads
                                                |
                                  distributed gradient finalization
                                                |
                                           optimizer step

The response groups execute serially. Their gradients accumulate into the same LoRA parameters; this is intended to keep the live response graph independent of the number of group members. The repository does not claim a quantitative memory reduction until a matched memory trace is recorded.

Resident state

glm52_prefix_state.py owns compact page state and its lifecycle. The stored components are model-specific:

• MLA latent key/value pages needed to reconstruct response attention.
• DSA index-key pages needed for sparse candidate selection.
• Per-layer metadata used to restore the exact structural prefix boundary.

The page manager records allocation, storage pointer, and PyTorch version state before response replay, then verifies that append/pop operations restore that structural witness. It is not a byte-content digest; raw storage writes are a backend contract violation, while numerical content is checked by parity gates. Prefix capture does not keep the original full-sequence autograd graph alive.

Global context-parallel DSA

Each CP rank initially scores candidates from its local prompt shard. Correct global sparse attention requires three additional steps:

1. gather rank-local candidate scores and global positions;
2. merge and deterministically select the global top-k set;
3. materialize the selected keys/values on the ranks that execute replay.

glm52_response_replay_adapter.py implements this candidate exchange and selected-key composition. The 32K forward parity result validates the composed selected log-probs. The newer 32K full-gradient comparison reaches 0.9999336 cosine and 0.0117/0.0116 relative L2; both cases pass the project gradient-parity acceptance contract.

Response-only update

For each GRPO group member, the worker integration:

1. restores the resident prefix boundary;
2. derives response inputs from the actual response tokens;
3. replays the decoder and chunked output head;
4. selects token-aligned log-probabilities;
5. applies the clipped GRPO objective and optional reference term;
6. backpropagates into configured LoRA targets;
7. restores the prefix metadata/storage/version boundary before the next group member.

After all group members, the worker calls the backend's finalize_model_grads surface and invokes the optimizer once. The current tree does not contain an independent finish_grad_sync call, so numerical parity is the required proof that the integration matches the selected backend schedule.

Online rollout and extended positions

glm52_online_grpo.py coordinates the external policy-sampler transaction: it exports the current LoRA, samples bounded real completions, computes DAPO math rewards, saves old log-probabilities, releases sampling actors, restores the Megatron actor, and submits the full rollout JSON for one finalized update. Sampling and training never occupy the same 32-GPU allocation concurrently. The accepted exact-2M transaction used the Tinker multi-node policy engine at TP8/PP4 for two eight-sample requests before restoring TP1/CP32/EP32 Megatron.

The released checkpoint remains native 1M. glm52_long_position.py provides a strict, opt-in GLM52_ROPE_SCALING_JSON boundary that applies the same experimental YaRN capacity configuration to policy sampling and Megatron. Without the environment setting, native model defaults are unchanged. The exact-2M online receipt validates this capacity path; it does not claim that untrained extended positions already have high task quality. The position contract state is ready_for_2m_training and remains on YaRN for the initial runs; alternative scaling is deferred until a trained model exposes a concrete position-dependent failure.

Direct Ray resource contract

build_direct_ray_plan keeps the distributed placement identity at the top-level node_ips, gpus_per_node, runtime_image, and image_resource fields. For a live digest, image_resource is the deterministic mint2m_image_<64hex> key derived from runtime_image. The placement-group bundles, coordinator spec, and each rank-actor spec bind both the target node:<ip> resource and this digest resource. RayRuntime.validate_cluster runs after ray.init but before placement-group creation and checks every target IP has one live Ray record, the requested GPU count, at least that many CPUs, its node resource, and the matching digest resource.

The four target Ray worker processes must already be running inside the deployment image at that digest with GPU devices attached. The direct Ray adapter's direct Ray actors do not receive a runtime_env or per-actor image override; actor processes inherit the cluster image. A review plan (runtime_image=None) remains useful for inspection and deterministic naming, but DirectRayDriver rejects it before creating the ownership journal or connecting to Ray.

Integration surfaces

The implementation is intentionally split by responsibility:

Surface	Responsibility
glm52_prefix_state.py	resident page ownership and lifecycle
glm52_prefix_capture.py / glm52_live_prefix_hooks.py	capture hooks and model boundary discovery
glm52_response_replay_adapter.py	global-CP sparse state and response replay
glm52_resident_prefix.py	engine-neutral GRPO math and dry-run contracts
glm52_chunked_lm_head.py	chunked output projection and loss
training_contract.py / training_session.py	explicit requests and ordered G2/optimizer lifecycle
training_backend.py	source-bound backend capability and cleanup contract
resident_replay_backend.py	transport-neutral resident replay, bounded gradient witnesses, and rank-port receipts
direct_ray.py / direct_ray_driver.py	pure placement plan, ownership journal, rollback, and recovery
ray_actors.py / ray_runtime.py	thin lazy-Ray wrappers; no HTTP control plane
doctor.py / validation_contract.py	prerequisite aggregation and ordered evidence identity
glm52/*_worker_patch.py	legacy private-runtime integration and audit export
runtime_patches/sitecustomize.py	legacy opt-in compatibility bootstrap
glm52_gradient_parity.py	conventional-versus-replay validation driver
glm52_online_grpo.py / glm52_validate_online_run.py	real policy rollout transaction and artifact acceptance
dapo_data_audit.py / dapo_math_reward.py	real-row validation and rule reward
glm52_long_position.py	reversible policy-sampler/Megatron position configuration
integrations/	model preparation and legacy runtime compatibility
glm52.cli	direct doctor, plan, train gate, and scoped cleanup
scripts/preflight_training.py / scripts/glm52_run.py	legacy profile execution pending migration

The worker and runtime patch files are predecessor integration adapters, not the external execution path. The repository-local megatron_rank_port.py, megatron_resident_runner.py, megatron_backend.py, and direct-Ray actors are the candidate external path. Candidate execution is enabled, while execution/redistribution authorization, the immutable runtime, approved inputs, and source-bound public validation remain independent gates. External source-unbound artifacts now record an exact-2M 32-rank online update and the accepted 32K gradient comparison. A direct_megatron_backend_unavailable result now means the backend is missing, incomplete, or explicitly disabled. Private HTTP/session/scheduler code is outside that boundary. See docs/upstream_integration.md.

Training evidence boundary

The direct runner binds a real model snapshot manifest and strict training fixture before creating Ray actors. Each rank reports the exact phase sequence and both planned and observed placement fields; the runtime checks node agreement, contiguous planned slots, and unique observed accelerator assignments. The driver persists those reports in artifacts/rank-sessions.json, records rank-local and driver cleanup, and emits artifacts/checkpoint-manifest.json. Receipt validation independently re-hashes the checkpoint tree and rejects symlinks or file drift. CPU/fake backends exercise these contracts but are not GPU or numerical-parity evidence.