[bebetterest]

Discussion: #49488

This draft exposes a reference implementation of configurable decoder layer execution from transformers_loop. It lets a supported model load its original checkpoint, execute layers according to a configurable plan, share the original parameters across repeated visits, and retain independent cache/state histories for each logical execution.

The immediate purpose is technical discussion and an inspectable implementation for researchers. I am not requesting that the full extension be merged in its current form. The design can remain external, become loop-specific model definitions, or provide a starting point for smaller interfaces agreed with maintainers.

AI assistance disclosure: Implementation, tests, documentation, and these submission drafts were developed with substantial code-agent assistance. The current first-time contributor restriction applies; this draft does not claim to satisfy the PR template's assertion that the code and description were not written by an LLM or code agent.

Execution model

For a four-layer decoder, [0, 1, 1, 2, 3] repeats only the second layer. [0, 1, 2, 3, 0, 1, 2, 3] repeats the complete stack as a block.

from transformers import RepeatRange

# model is a supported model loaded using the fork or this implementation.
model.set_layer_execution_plan(repeats=[RepeatRange(1, 2, total_passes=2)])
model.set_layer_execution_plan(repeats=[RepeatRange(1, 4, total_passes=2)])
model.set_layer_execution_plan([0, 1, 1, *range(2, model.config.get_text_config().num_hidden_layers)])
model.set_layer_execution_plan(None)

Each call replaces the plan. Indices are zero-based, ranges are half-open, and the pass count includes the first visit. Plans must change between complete training steps; existing caches must be discarded after a change.

• Physical source layers retain their registered modules, parameters, optimizer registration, and checkpoint weight keys.
• Logical execution positions select independent cache histories. KV, convolution, recurrent, and custom state are routed through execution-local cache/state views.
• Embeddings and stack input transforms run once, followed by the planned layers and then final normalization/output transforms. Token positions advance once per forward.
• Autograd accumulates contributions from every repeated visit. Disabling the plan restores native execution with the current weights.
• Configuration serialization preserves the execution order and optional native KV dependency bindings. A plan can also be specified before loading the original weights.

Changes included

Area	Implementation
layer_execution/	Immutable plans, repeat ranges, adapter registry, executor, execution-local caches/states, and backend integrations
Model adapters	Llama; dense Qwen3.5 text; Gemma3n text with native KV dependencies and AltUp streams; RecurrentGemma with externalized persistent state
Configuration and loading	Opt-in plan persistence and restoration while retaining physical weight ownership
Generation and masking	Loop-aware cache creation, supported static/offloaded caches, assisted-generation state handling, and selected compiled paths
Continuous and paged batching	Portable scheduling and supported attention-only paged execution with logical cache slots
Distributed execution	Native DDP/FSDP2/TP integration and standalone sequential PP with source-layer ownership and state transfer
Trainer and DeepSpeed	Plan-aware training/checkpoint paths, adapter resume fixes, and ZeRO-3 shared-gradient coordination
Tests and docs	Plan/cache/gradient oracles, save/reload, PEFT/quantization, distributed regressions, and API documentation

This reference diff contains 42 files, with 7,114 added and 64 removed lines, against upstream 536ecc007387a50e77603bb5d92100e9b07514cc. Fork branding, the root loop guide, and the fork-maintenance document are excluded from the proposed diff. Their public links remain available as supplementary usage and validation records.

The execution protocol targets causal decoder-only stacks whose prepared state streams preserve their shapes and structure across layers. Other models need an architecture adapter; a matching hidden width alone does not establish support. GPT-2, OPT, and Mamba are extension examples/tests, not automatically registered built-ins.

Recorded validation

The source/test contents are taken from fork commit cf2d29f84c9b9449ac639a687737b9d957bc9fc2, synchronized with upstream 536ecc007387a50e77603bb5d92100e9b07514cc. The following are previously recorded results for those contents, not claims about this PR's CI.

Environment or check	Recorded result
Local CPU, Python 3.11.15 / PyTorch 2.14.1	123 loop tests passed / 88 skipped; 10 Trainer checkpoint tests passed; 31 affected upstream tests passed / 1 skipped
Two H20 GPUs on one physical host, Python 3.12.3 / PyTorch 2.12.1+cu130	204 unique native loop tests passed / 7 skipped, plus 8 passing subtests
Selected single-GPU TRL checks after synchronization	3 SFT/DPO/GRPO cases with independent loss/gradient oracles and resume controls
Selected two-GPU TRL regression matrix after synchronization	25 configurations, 50 successful process phases, and 100 checked rank reports
Independent resumed/uninterrupted tensor audit	22 parameter-file pairs matched exactly across 11 tiny-model control cases and two ranks

The GPU environment used unmodified TRL 1.15.0, PEFT 0.21.2, Accelerate 1.15.0, DeepSpeed 0.19.7, and bitsandbytes 0.50.2. The selected matrix includes DDP, ZeRO-2/3, full-parameter training, LoRA, QLoRA, and original Qwen/Qwen3.5-9B weights at revision c202236235762e1c871ad0ccb60c8ee5ba337b9a. Original full-parameter 9B checks use SGD; tiny and adapter runs use AdamW.

The pinned maintenance record provides the selected matrix, dependency versions, skips, resolved bootstrap/environment failures, and interrupted attempts excluded from passing counts. Earlier full matrices belong to a separate fixed baseline and are not counted as reruns after synchronization.

Relevant local commands are:

PYTHONPATH=src HF_HUB_OFFLINE=1 python -m pytest tests/test_layer_execution*.py -q
make check-repo
make style

GPU and optional-backend cases require the corresponding hardware and dependencies. The recorded GPU results apply only while the submission branch's source/test contents match the pinned implementation. Submission-branch checks are recorded separately:

Checks run on the submission contents on 2026-10-11 (commit db1a3aeb03551d5e84422f1851947648c6b93c88):

• Verified that all 4,776 Python files under src/ and tests/ match the pinned fork contents, and that the committed 42-file diff matches the reviewed patch.
• Loop suite: 123 passed, 88 skipped on local CPU. Hardware and optional-backend skips are retained.
• Selected Trainer checkpoint/resume cases: 14 unique cases passed, 2 skipped after resolving one offline fixture-loading failure. The first attempt was 13 passed / 2 skipped / 1 failed before training; the failed case passed after caching the public tiny-Llama fixture at revision 9fb191250dd56d0ba7ec9785a025ed29c03d5998. No runtime source change was needed.
• Selected affected upstream configuration, loading, and masking regressions: 30 passed, 1 skipped.
• make check-repo succeeded. A separate --no-cache run completed all 26 checker statuses successfully; the reviewer-assignment check skipped its work because the optional shared resolver is not installed. make style completed all 5 checks successfully and left the committed diff unchanged.
• Attempted utils/tests_fetcher.py --diff_with_last_commit. Its reverse-dependency discovery failed on attention constants interpreted as module paths under cache_allocators/. The same class of failure was reproduced on unchanged upstream 536ecc007387a50e77603bb5d92100e9b07514cc; no affected-test list or full-suite passing result is claimed.
• This submission pass did not rerun the historical GPU matrix. Upstream PR CI results, when available, are separate from these local checks.
• Upstream CI after creation: the automatic security gate reported one LOW-severity B603 finding at tests/test_layer_execution_multinode.py:112 for subprocess.Popen. This test uses a sys.executable argument list, without a shell, to start two local torchrun launchers for same-host simulation. The gate run and bot comment show that manual maintainer review/approval is required; full upstream CI has not been automatically approved. The test source is retained for review.

Boundaries and design tradeoffs

• vLLM is not integrated or runtime-validated, including model_impl="transformers" and TRL server/colocate modes. GRPO uses Hugging Face generation with use_vllm=False.
• Physical multi-host execution remains unverified. Two independent launchers on one host do not establish real multi-host behavior.
• Native FSDP2/TP/PP checks do not establish their TRL combinations. Native load-time TP/FSDP2 requires model-only saving; optimizer-state resume is unsupported on that path.
• PP scheduling is sequential. Overlapped GPipe/1F1B and combined parallel layouts are outside the integrated scope.
• Optional quantized-cache backends, optimized recurrent kernels, other dependency versions, ZeRO-3 initialization during loading, and CPU/NVMe optimizer offload need their own validation.
• This implementation binds an execution forward to supported decoder instances. ZeRO-3 coordination uses version-sensitive gradient-reduction APIs and a PyTorch internal autograd callback. These are explicit maintenance tradeoffs for discussion.
• Short integration checks establish execution/training mechanics, not improved model quality, convergence, production throughput, or long-run reliability.
• Adapter-only saving does not preserve a plan set only in memory. Save and reload the planned base checkpoint when automatic reference-model construction or adapter reload must use the same plan.

Related interfaces include KV sharing and heterogeneous layer configuration. The linked issue asks whether this should remain external, use distinct model definitions, or be reduced to a smaller core interface. Maintainer approval for such an upstream design has not yet been obtained.

[github-actions]

Thank you for your contribution 🤗!

CI Security Gate — automatic approval blocked

This PR was not automatically approved for CI because the security gate failed.

Possible reasons:

• The PR touches 50 or more files — only PRs with fewer than 50 changed files are automatically approved
• A changed file is outside the allowed directories (src/, tests/, docs/, utils/), has a disallowed extension (only .py, .txt, .md permitted outside tests/ and docs/), or is not .md/.yml inside docs/ — this covers files the PR deletes or renames, not only the ones it edits
• A new high-severity security issue was detected in the changed Python files (Bandit check)
• The PR touches a path this repository protects from untrusted PRs, such as the file that decides who reviews it — a maintainer must make that change in a separate PR

See the workflow run for the exact violations.

A maintainer can review and manually approve CI if a finding is a false positive.