Hugging Face's logo

viai957
/
align

lean4
theorem-proving
post-training
grpo
code
Model card Files
xet
 Community
align / code /docs /ARCHITECTURE_V2_PROPOSAL.md
viai957's picture
viai957
Add files using upload-large-folder tool
d909bd8 verified
preview code
|
Raw History Blame Contribute Delete
82.6 kB

align architecture v2: proposal

Status: proposal only. No code was moved or edited to write it. Date: 2026-10-14. Scope: /apps/development/rlhf/align at HEAD 1c483cc, compared with reference-repos/AgileRL at 4f07424 (Apache-2.0).

How to read the evidence labels used below.

  • [read] read in source; the path and line are given and can be re-checked with sed -n.
  • [calc] arithmetic from numbers that are themselves labelled. The arithmetic is shown or reproducible in five lines.
  • [cpu-probe] a small CPU run done for this document (the probe is described where it is used; nothing was left in the repo).
  • [assumed] a number taken from a datasheet or a typical value, not measured on this box. Every one is listed in 5.1.
  • [paper] arXiv id plus section or equation, checked against the text fetched into the hyperresearch vault.

Independent verification pass (2026-10-14, builder A2): the freeze closure, the layer assignment, the shim mechanics, the source-pinning tests and the Section 5 arithmetic were re-derived from the repo; corrections are folded into the text below and marked [A2].

Nothing here was measured on a GPU (the P2 chain owns all four). Section 5 is therefore a model with stated assumptions and a measurement plan, not a benchmark. The CPU test suite was not run for this document either; the three docs that quote a test count disagree (1089 in DEPENDENCY_POLICY, 1107 in REFERENCE_GAP_ANALYSIS, 1608 in README), so step M0 records the real baseline first.

0. Summary of decisions

  1. Adopt from AgileRL: interfaces and contracts, not structure. Seven ideas are worth taking (1.3): a typed, validated experience container; algorithm variants as a thin switch over one implementation (CISPO and GSPO are GRPO with a fixed loss_type); importance-sampling level and advantage granularity as two independent axes; a text-env protocol separated from a token-level rollout harness; an async collector; a documented-constructor AST test; a shard-runtime seam for dense versus FSDP2.
  2. Do not adopt its structure or its dependencies. The 6,398-line agilerl/algorithms/core/base.py makes the algorithm own generation, the optimizer step, vLLM colocation, FSDP and LoRA plumbing. align's Trainer / TrainModule split is already better. accelerate, transformers, peft, vLLM, liger-kernel, pydantic, hydra, wandb, dill-pickled checkpoints and the evolutionary HPO stack are all out (1.4).
  3. The one structural gap in align is the environment. GrpoBatchSource is hard-wired to {problem, answer} rows plus an outcome verifier; the multi-turn Lean path is a subclass that re-implements bind. A single Env protocol removes the fork.
  4. Layering is currently violated in 5 places by lazy imports (2.1; [A2] one is an artefact of putting eval at L5, two are missed ones). They are cheap to fix and a 40-line AST test keeps them fixed.
  5. Migration is facade-first. New packages (algorithms/, envs/, recipes/, optim/) are added as thin layers over the old modules. Old modules move only after the P2 chain exits, with identity-preserving sys.modules shims (verified in 4.2). [A2] 105 of the 129 modules in align/ are in the import closure of the running P2 scripts and of the GPU-validation steps the chain scripts launch; they cannot be touched for now (4.1).
  6. Kernels: only one custom Triton kernel pays a lot, and it is not one of the three the task listed first.
    • Paged-attention decode (TritonPagedAttention, an interface with no kernel today): about 2x to 4x on decode steps at batch 6 to 64 [calc, 5.4-C; [A2] rows that exceed the KV capacity of the card were removed]. This is the big one, because rollout generation dominates an RL iteration (5.2).
    • A fused head + policy-gradient loss pays about 4x on the head (about 13% of an RL training step) but the gain comes from precision path and GEMM count (4 fp32 GEMMs to 2 bf16 GEMMs), not from fusing elementwise ops. The elementwise part is about 4% of today's head time and 10% to 16% of the fused head's time [A2]. Do it as a chunked autograd.Function first; Triton is optional there (5.4-A, 5.4-B).
    • Fusing only the per-token CISPO/GRPO loss arithmetic is worth nothing (64 KB tensors).

1. AgileRL: what to adopt and what not to

1.1 What AgileRL is

An evolutionary-HPO RL library (populations, mutation, tournament selection) that grew an LLM fine-tuning half. Size by directory [read, find ... | xargs cat | wc -l]:

dir lines relevance to align
algorithms/ 27,282 core/base.py 6,398; grpo.py 2,471; ppo_llm.py 1,415; reinforce_llm.py 1,125; dpo.py 692; sft.py 410; cispo.py 26; gspo.py 23. The rest is classic RL (DQN, TD3, MADDPG, bandits)
utils/ 17,179 llm_utils.py 3,035, algo_utils.py 2,771, probe envs, mutation utils
modules/ + networks/ + wrappers/ + vector/ + hpo/ + population.py 15,612 evolvable networks and populations; none of it applies
llm_envs/ + training/ + strategies/ + builders/ + models/ 11,496 the environment and trainer side; partly relevant
distributed/ 3,112 FSDP2 runtime, process helpers
algorithms/core/llm_ops/ fused ops fused_logprobs.py, fused_loss.py (Liger), fused_lora.py, moe_lora.py

Roughly 40% of the code is classic-RL or evolution machinery that an LLM post-training library does not need. That is the first reason not to copy its layout wholesale.

1.2 How its LLM algorithms are built

Class tree [read]. EvolvableAlgorithm (algorithms/core/base.py:436, metaclass RegistryMeta at :312 runs _registry_init after every constructor) holds mutation registry, metrics, accelerator and device. LLMAlgorithm (base.py:2687, about 3,700 lines of the file) adds: actor/reference/critic adapters, vLLM colocation and weight sync (_sync_actor_to_vllm :5466, _generate_with_vllm_colocate :5484), FSDP shard runtime, fused logprob forward (_fused_forward :4844, _fused_packed_forward :4915), LoRA attach (_attach_lora_adapters :4341), backward and grad clip (_backward_pass :5283), checkpointing (save_checkpoint :3157).

Algorithms. GRPO(LLMAlgorithm) (grpo.py:170; ctor :398 with about 50 parameters; get_action :541 generates; learn :634 runs the epoch/minibatch loop, loss, backward and step). CISPO(GRPO) (cispo.py:15) and GSPO(GRPO) (gspo.py:15) are about 25 lines each: they call super().__init__(..., loss_type="cispo"). The objective is chosen in _setup_objective (grpo.py:1071): loss_type in {grpo, gspo, cispo}, importance_sampling_level in {token, turn, trajectory} (gspo forces trajectory), advantage_granularity separately, loss_norm in {micro_batch, accumulation_window}. DPO (dpo.py:44) and SFT (sft.py:42) follow the same get_action / learn(experiences) shape with a typed experiences argument.

ILQL and BC_LM are a second generation. ilql.py:42 is class ILQL(nn.Module) of 2,256 lines, a port of an older offline-RL codebase with its own Language_Environment / RL_Dataset types, and import wandb at module top (ilql.py:15). It does not use LLMAlgorithm. So AgileRL itself carries two incompatible LLM-algorithm architectures. That is the cost of growing a base class instead of defining a contract first.

Experience container [read, components/llm_rollout_data.py]. Trajectory (frozen pydantic model: token_ids (1,T), action_masks/turn_ids (1,T-1), rewards per turn, optional sampling_logps, pixel_values) with a model_validator that fails when per-token fields are not T-1 long. RolloutGroup requires exactly group_size trajectories. They are packed into LLMExperienceBatch. The label-aligned T-1 convention is the same one align states in rl/types.py.

Environments [read]. protocols.py:602 TextEnvProtocol (reset(...), step(action_text), system_prompt) and :624 EnvClientProtocol (Gym five-tuple, dataset_size, tools, rubric_components, eval_mode()). RolloutHarness (llm_envs/harness.py:44, 857 lines) does tokenisation, chat-template boundaries and the turn loop over a client. RolloutCollector (collector.py:28) runs many harnesses; TaskAssigner (task_assigner.py:27) assigns dataset rows to slots; AsyncBatchCollector (async_collector.py:24) overlaps collection with learning; rubrics.py composes reward components.

Training loop and config [read]. training/llm/rollout.py:100 train_llm_rollout(pop, max_turns, env_factory, ...). Config is YAML under configs/training/llm_finetuning/ (cispo, dpo, grpo, gspo, ppo_llm, reinforce_llm, sft) parsed into pydantic *Spec models (models/, builders/llm.py:46 LLMBuilder).

Checkpointing [read, base.py:3157-3250]. A directory with PEFT adapter dirs plus attributes.pt, a pickle written with dill (base.py:33, :1366, :1381). FSDP2 state is gathered to full tensors before saving, so checkpoints are rank-count independent. Only rank 0 writes.

Distributed [read]. accelerate for process setup, plus a native FSDP2 layer: distributed/runtime.py defines a shard runtime (prepare_actor, export_model_state, export_optimizer_state) with dense and FSDP2 implementations.

Tests and docs [read]. tests/ mirrors the package (test_algorithms/test_llms, test_llm_ops, test_llm_envs, ...). test_algorithms/test_param_docstrings.py is an AST test: every constructor parameter of an audited class has a :param: entry and no stale entries remain. Sphinx docs under docs/llm_finetuning/ (batch sizing, fused logprobs, checkpoints, environments).

Convergent design worth noting. AgileRL's fused_logprobs.py (fused_linear_logprobs_chunked :135, FusedLinearLogProbsFunction :199) is the same chunked head + log-softmax + gradient-checkpoint idea as align's kernels/fused_ce.py and rl/cispo.py::fp32_head_logprobs. AgileRL's own GRPO docstring says its Liger GRPO kernel "shows no speedup over AgileRL's already memory-bounded standard path" (grpo.py, use_liger_loss parameter doc). That matches the roofline result in 5.4-A: chunking removes the memory problem, and a fused elementwise kernel adds almost nothing on top.

1.3 Ideas to adopt

# Idea AgileRL reference align today Concrete action
A1 Typed, validated experience batch with a shape-contract validator components/llm_rollout_data.py Trajectory, RolloutGroup rl/types.py RolloutBatch.__post_init__ (:51) already validates; TokenAdvRolloutBatch lives in rl/grpo_loss.py:301; DpoBatch in train/dpo_stage.py:818; SFT batch is an untyped dict from collate_sft One module align/batch.py holding SFTBatch, PreferenceBatch, RolloutBatch, TokenAdvRolloutBatch, PPORolloutBatch, each with validate(). Dataclasses, not pydantic
A2 Variants as a thin switch, not copies cispo.py / gspo.py (25 lines each) GRPOLossType.{dapo,cispo} in rl/grpo_loss.py:41 already does it for the loss; GSPO absent algorithms/grpo/ owns one implementation; algorithms/cispo.py and dapo.py are presets (2.3). Add the trajectory-level ratio as a third option only when needed (GSPO, arXiv:2507.18071 Sec. 4.1 Eq. 5-7: s_i = exp(mean_t log(pi/pi_old)))
A3 Two independent axes: importance-sampling level and advantage granularity grpo.py:1071 _setup_objective, :1493 _resolve_advantage_granularity rl/multiturn_source.py computes per-turn advantages, expand_advantages (grpo_loss.py:234) broadcasts; ratio level is token only Name both axes in GRPOConfig (ratio_level, advantage_level) and refuse the mismatched pair that AgileRL only warns about (grpo.py warns on turn advantage with trajectory ratio)
A4 Text-env protocol separate from the token-level harness protocols.py:602, llm_envs/harness.py:44 rl/envs/lean_multiturn.py fuses both; LeanMultiTurnEnv.run_group(specs, generate) already takes the generator as a callable, which is the right seam align/envs/ (2.4): Env protocol over text, one Harness for tokenisation and per-turn masks, MathEnv, LeanEnv, CodeEnv
A5 Async collection overlapping learning llm_envs/async_collector.py:24 rollout is synchronous (rl/batch_source.py _produce_step); REFERENCE_GAP_ANALYSIS 1.3 already lists this Defer. Design the RolloutBackend.submit/collect protocol (already present in rollout/protocol.py:95) so a one-step-stale pipeline is a source-level change, not an engine change
A6 Documented-constructor AST test tests/test_algorithms/test_param_docstrings.py align documents fields as #: comments on dataclasses (config.py) but nothing checks it tests/contract/test_config_docs.py: every field of every registered algorithm Config has a doc comment. Stdlib ast, runs in the firewall lane
A7 Shard-runtime seam (dense vs FSDP2 vs TP) distributed/runtime.py TrainModule.full_model_state_dict, optim_state_dict, load_* already are this seam (train/train_module.py), plus model/tp.py for inference TP Name the seam ShardRuntime in align/protocols.py and make TrainModule and TPPagedEngine implement it; no code change to the Trainer

Also adopted as a habit rather than code: AgileRL keeps fused ops in their own package with their own tests (algorithms/core/llm_ops/, tests/test_algorithms/test_llm_ops/). align already does this in kernels/.

1.4 Ideas not to adopt

Idea AgileRL reference Why not
A 6,398-line base class that owns generation, vLLM colocation, FSDP, LoRA and the optimizer step algorithms/core/base.py The algorithm cannot be tested without those. align's TrainModule owns model+loss+optimizer, Trainer owns loop and state, the batch source owns rollouts. Each is testable on CPU alone (1,000+ CPU tests prove it)
get_action generates inside the algorithm grpo.py:541 Couples the algorithm to HF generate or colocated vLLM. align's RolloutBackend (rollout/protocol.py:95) keeps the engine replaceable (native, paged, spool/vLLM)
Replacing lm_head with a pass-through HiddenStateHead (documented as nn.Identity) by monkeypatch to get hidden states base.py:6314 _patch_lm_head_to_identity Mutates the module tree under a context manager. align passes hidden_consumer(hidden, lm_head_weight) into forward (model/qwen3.py:467, mistral.py:597): explicit, and it fails loudly on a model without the kwarg (rl/cispo.py fp32_head_forward)
accelerate, transformers, peft, bitsandbytes, liger-kernel, vLLM as hard or core dependencies pyproject.toml dependencies, cpu-llm extra Violates docs/DEPENDENCY_POLICY.md and tests/test_dependency_firewall.py. align already has the replacements (own LoRA, own 8-bit AdamW, own tokenizer, own paged engine)
pydantic spec models and hydra/omegaconf models/, builders/ align/config.py is dataclass + tomllib with struct mode and typed dotted overrides, no dependency. Adopt only the idea that the algorithm's config is its own dataclass
dill-pickled attributes.pt as the checkpoint container base.py:33, :1366 Pickle loads arbitrary code. align reads resume.pt with weights_only=True and has fsync, partial-staging and completeness checks (train/trainer.py) AgileRL lacks
Evolutionary HPO: population, mutation, tournament, EvolvableModule networks hpo/, population.py, modules/, networks/ Irrelevant to single-run post-training, and it shapes the base class (MutationRegistry, _registry_init metaclass hook)
The older nn.Module ILQL/BC_LM generation algorithms/ilql.py, bc_lm.py Shows the failure mode of not fixing the contract first (1.2)
Arena/cloud manifests models/manifest.py, agilerl-arena Not a use case here
Warn-and-continue on conflicting options many warnings.warn in _setup_objective align's convention is a launch refusal (assert_a3_conformant, assert_p4_cispo_conformant). Keep refusing

1.5 Where align is already ahead (keep as is)

  • The Trainer / TrainModule split with extra_state() hooks and a documented resume flow (docs/ARCHITECTURE.md, "Trainer / TrainModule split").
  • The control plane: coordinated stop, MAX-reduce decisions, exit codes 80 to 83, HALT.json, supervisor (align/exit_codes.py, utils/supervisor.py).
  • Host safety: preflight, heavy-job flock, D-state gate (utils/guard.py, utils/host_state_gate.py).
  • The dependency firewall as an executable test.
  • Launch-refusal gates per stage with a dated rationale in each docstring.
  • Bit-invariance as a first-class engine mode (TorchPagedAttention(mode="row")).

None of this exists in AgileRL, and none of it should be traded for its structure.

Licence note: AgileRL is Apache-2.0. Adopting ideas needs no attribution; copying code needs the header and a NOTICE entry. This proposal assumes ideas only.

2. Target package layout

2.1 Layer rules, and where the current code breaks them

Imports flow downward only. A layer may import its own layer only through an explicit allow-list in the test.

L0  root modules: config, exceptions, exit_codes, registry(banner), protocols (new), batch (new)
    runtime/                     (was utils/)
L1  kernels, tokenizer, model, optim (new), functional (new: pure tensor/scalar math, no align imports)
L2  verify, data, rollout, inference, envs (new)
L2b eval             (imports only L1/L2: data, verify, inference, model, tokenizer; measured, see below)
L3  train            (the engine: Trainer, TrainModule, checkpoint formats, FSDP, stage gates and smoke)
L4  algorithms (new) (losses, batch contracts, modules, presets; subclass TrainModule)
L5  recipes (new), cli (new)

[A2] Why eval is L2b and not L5. The first draft put eval at L5 and then listed train -> eval as an upward import. An AST scan of align/eval/*.py shows its only align imports are config, data.schema, data.tokenization, inference.engine, model.*, tokenizer, verify and itself, so eval sits below train. With eval at L5 the proposed fixes created two new violations (moving train/stage_gates.py to recipes/ makes dpo_stage, now L4, import L5; moving ref_cache to algorithms/dpo/ leaves ref_cache -> eval.metrics.wilson_interval, L4 to L5) and moved a 955-line, 12-import-site, frozen file for no gain. With eval at L2b, stage_gates.py, stage_smoke.py and ref_cache.py need no layering fix. The align/__init__.py lazy public API is exempt from the layer rule (it names L2 modules from L0 through __getattr__).

Edges measured with an AST scan of every import including lazy ones [read, script run for this document]:

Violation Evidence Fix
train imports eval [A2] not a violation once eval is L2b (see above) train/stage_smoke.py:369,370,453,463, train/ref_cache.py:594, train/stage_gates.py:351,685 none: stage_gates, stage_smoke stay in train/; ref_cache moves to algorithms/dpo/ and may import eval
train imports ppo (upward) train/qwen3_stage.py:171,189,209,224,258 (ppo.lora) and :333 (ppo.trainer.AdapterProbeView) [A2: the first draft listed only 171,189,209 and missed the AdapterProbeView edge] LoRA is a model concern: ppo/lora.py to model/lora.py; AdapterProbeView is a parameter view for train.probes.TrunkMovementProbe, so it moves next to it, ppo/trainer.py::AdapterProbeView to train/probes.py
[A2] new: rl/envs/lean_multiturn.py (becomes envs/, L2) imports rl.advantages and rl.rewards (would be algorithms/, L4) lean_multiturn.py:100,101 (compute_group_advantages, soft_overlong_punishment) rl/advantages.py, rl/rewards.py, rl/ops.py, ppo/gae.py have no align imports (checked); they become functional/{advantages,rewards,ops,gae}.py at L1 instead of algorithms/common/
data imports eval while eval imports data (cycle) data/outcome_agreement.py:275,306 imports eval.verifier; eval has 3 edges into data Outcome/OutcomeVerifier to verify/outcome.py; outcome_agreement.py to eval/
verify imports eval while eval imports verify (cycle) verify/lean/reward.py:116 imports eval.verifier.Outcome same move
The library loads code from scripts/ and extras/ by file path rl/envs/lean_multiturn.py:108-123 (_load_by_path("...", "extras/vllm_rollout/lean_prompts.py"), "scripts/p2/p2_common.py") Vendor the prompt module as envs/lean/prompts.py with a sha256 identity test against the extras/ copy; extract the one function used from p2_common (displayed_statement_matches)
Hard-coded absolute paths in library code config.py:40,335,339, kernels/gate.py:36, kernels/_toolchain.py:42, utils/guard.py:41,42,43 (PROJECT_ROOT = "/apps/development/rlhf/align"); eval/golden.py:84,123,128,129,374; verify/lean/toolchain.py:39; rollout/client.py:135; data/hub.py:54, data/sources.py:37, data/ingest_parquet.py:36; train/dpo_stage.py:142,149 ([A2] the last three groups were missing; config.py:335,339 too) M8; blockers for open-sourcing. Editing eval/golden.py changes eval_source_hash (4.5)

The enforcing test (tests/test_layering.py, about 40 lines) builds the same graph, maps each module to a layer from a table in the test, and fails on an upward edge or on a cycle between two same-layer packages. [A2] It scans only real modules: files listed in the shim manifest (4.2) are skipped, and the legacy packages ppo, rl, rl.envs, utils are treated as aliases of their targets, otherwise every shim would be reported as an upward edge. It starts in report-only mode (M0) so the current violations are recorded, then becomes strict as each is fixed.

2.2 Tree 

align/
  __init__.py                 lazy public API (existing pattern; new names added in M7)
  config.py exceptions.py exit_codes.py registry.py            unchanged location (L0)
  protocols.py                NEW  Algorithm, Env, RolloutBackend, Verifier, ShardRuntime, WeightFormat, LossAdapter
  batch.py                    NEW  SFTBatch, PreferenceBatch, RolloutBatch, TokenAdvRolloutBatch, PPORolloutBatch
                              ([A2] both files import torch only under TYPE_CHECKING or inside functions; neither is imported by
                              `align/__init__`, and `test_public_api`'s heavy-module list gains `align.batch`, `align.protocols`)
  runtime/                    dist.py guard.py host_state_gate.py supervisor.py tracking.py   (from utils/)
                              __init__.py stays empty: `import align.runtime` must not import torch (tests/test_public_api.py)
  functional/                 NEW (L1)  advantages.py rewards.py ops.py gae.py kl.py   pure math, no align imports
  kernels/                    unchanged + fused_policy_loss.py (NEW), paged_attention.py (NEW, Triton)
  tokenizer/                  unchanged
  model/                      unchanged + lora.py (from ppo/lora.py), heads.py (from ppo/value_head.py),
                              build.py (arch registry: mistral | qwen3 | olmoe)
  optim/                      adamw8bit.py quant8.py offload_adamw.py factory.py (make_optimizer)
  verify/                     unchanged + outcome.py (from eval/verifier.py), code/ (NEW)
  data/                       unchanged (minus outcome_agreement.py)
  rollout/                    unchanged + backends.py (NEW registry: native | paged | tp | spool)
  inference/                  unchanged (eval-time contiguous-KV engine; candidate to be replaced by PagedEngine later)
  envs/                       NEW  base.py harness.py math.py lean/{env.py,prompts.py} code.py
  train/                      engine: trainer.py train_module.py fsdp.py probes.py (+AdapterProbeView) sharded_ckpt.py ckpt_format.py
                              moe_module.py stage_gates.py stage_smoke.py   ([A2] the last two stay)
  algorithms/                 NEW  base.py registry.py
                              sft.py  dpo/  ppo/  grpo/  cispo.py  dapo.py  expert_iteration.py
  eval/                       unchanged + outcome_agreement.py (M7 only: see 4.5, changes eval_source_hash)
  recipes/                    NEW  sft.py dpo.py grpo.py ppo.py expert_iteration.py
  cli.py                      NEW  python -m align {sft,dpo,grpo,ppo,ei,eval}

Not moved on purpose: config.py, exceptions.py, exit_codes.py, registry.py (55 test import sites for config.py alone, all in the freeze closure, and nothing is gained by a core/ directory that the layer table does not already give); model/, kernels/, tokenizer/, rollout/, inference/, data/ names (97 test imports use align.model alone). Tests (121 files) stay flat in tests/; a new tests/contract/ holds the contract suites (2.3).

2.3 The algorithm contract

Today the pieces of an algorithm are spread over four places with no common name: a loss adapter (model, batch) -> LossOutput (train/train_module.py LossAdapter), a stage config parsed by its own loader (rl/stage_config.py:181, ppo/stage_config.py:313, train/dpo_stage.py:285), a launch-refusal function (assert_a3_conformant at rl/grpo_loss.py:337, assert_p4_cispo_conformant at rl/cispo.py:295), and a TrainModule subclass. The contract names them.

# align/protocols.py  (sketch)
@dataclass(frozen=True)
class Needs:
    reference: Literal["none", "adapters_off", "cached_logprobs", "model"]
    rollouts: bool
    critic: bool

class Algorithm(Protocol):
    name: ClassVar[str]
    Config: ClassVar[type]            # frozen dataclass, struct-mode loadable by align.config
    Batch: ClassVar[type]             # batch contract in align/batch.py, with .validate()
    needs: ClassVar[Needs]

    def __init__(self, cfg: "Config"): ...
    def validate(self, run: RunConfig, world_size: int) -> list[str]:
        """Launch refusals. Empty list = may launch. Replaces assert_a3_conformant et al."""
    def loss_adapter(self) -> LossAdapter:                      # (model, batch) -> LossOutput
    def train_module(self, model, optimizer, run: RunConfig) -> TrainModule: ...
    def make_source(self, data_cfg, *, env=None, engine=None, model=None, tokenizer=None): ...
    def extra_state(self) -> dict: ...                          # rides resume.pt["extra_state"] (exists)
    def load_extra_state(self, state: dict) -> None: ...
    def banner(self) -> dict:                                   # keys registered with align.registry

Rules the contract encodes, each already true somewhere in align and now stated once: train_batch is exactly (loss / grad_accum).backward(); per-token tensors are label-aligned [B, T-1]; extra state loads with weights_only=True; undefined metrics are NaN, never 0.0.

Mapping of the seven algorithms (current locations [read]):

Algorithm Loss today Batch contract today Config today Extra state Needs
sft make_sft_loss_adapter train/sft_loss.py:54 over fused_linear_cross_entropy (kernels/fused_ce.py:259) dict from data.loader.collate_sft RunConfig sections only none none
dpo make_dpo_loss_adapter train/dpo_loss.py:115 over fused_dpo_loss (kernels/fused_dpo.py:545) DpoBatch train/dpo_stage.py:818 DpoStageConfig :118 tripwire response (:680) cached_logprobs (train/ref_cache.py)
ppo ppo_loss ppo/loss.py:83, value_warmup_loss :163, PpoTrainModule ppo/trainer.py PPORollout ppo/rollout.py:40 PPOStageConfig ppo/stage_config.py:110 ratio tripwire (ppo/health.py) adapters_off reference, critic, rollouts
grpo compute_grpo_loss rl/grpo_loss.py:179, GRPOTrainModule rl/grpo_module.py:53 RolloutBatch rl/types.py:21, TokenAdvRolloutBatch rl/grpo_loss.py:301 GrpoStageConfig rl/stage_config.py:92, GRPOLossConfig rl/grpo_loss.py:47 CollapseTripwire state, GrpoBatchSource cursor (see "Still open" in DEPENDENCY_POLICY) rollouts; reference only if beta > 0
dapo GRPOLossType.dapo branch (grpo_loss.py:194) + dead-group filter (rl/advantages.py:50) + soft overlong (rl/rewards.py:22) as grpo as grpo as grpo as grpo
cispo cispo_policy_loss rl/cispo.py:110, fp32_head_forward :255 as grpo GRPOLossConfig(loss_fn=cispo) + assert_p4_cispo_conformant as grpo as grpo, beta = 0
expert_iteration inner loss is sft; outer is an 11-stage driver scripts/p2/expert_iteration.py (850 lines) with chain-hash markers stage markers (RUN/round_NN/stages/<stage>.done.json) argparse in the script round manifests none (subprocess stages)

Two honest wrinkles in this table.

  • "DAPO" in align is not DAPO's objective. GRPOLossType.dapo is the clip-higher surrogate (DAPO Sec. 3.1, arXiv:2503.14476, eps_high = 0.272 in grpo_loss.py:54), but the A3 gate refuses DAPO's own token-mean denominator (grpo_loss.py:56 comment "A3 requires a float (Dr.GRPO)"; refusal text at :356, "not the DAPO token-mean 'token' default") in favour of a constant denominator and centred advantages without std division (Dr. GRPO, arXiv:2503.20783, Sec. 3.1 on the length and std normalisation biases). Dynamic sampling (DAPO Sec. 3.2) is the dead-group filter; soft overlong punishment (DAPO Sec. 3.4) is rl/rewards.py:22. The presets should say what they are: dapo (upstream-faithful, token mean, for comparison), a3 (the pre-registered clip-higher + constant denominator), cispo_p4 (CISPO, arXiv:2506.13585 Sec. 3.1 Eq. 3-5, no lower IS bound, beta = 0, fp32 head per Sec. 3.2). The A3 and P4 gates stay as validate() and keep refusing.
  • expert_iteration is an outer loop, not a per-step loss. The contract has two kinds: a LossAlgorithm (everything above except this one) and an OuterLoop (a stage graph whose stages are subprocesses of other recipes). Do not force it into loss_adapter. It implements validate, banner and a run(stage_graph); its stages call recipes.sft, recipes.eval.

Contract test (tests/contract/test_algorithm_contract.py, parametrised over the registry): (1) Config round-trips through dumps_toml and load_config with struct mode; (2) loss_adapter() returns a finite LossOutput on a tiny fp64 model and synthetic batch; (3) the gradient equals an eager fp64 reference to a stated tolerance (the "oracle tier" pattern already used for model parity); (4) extra_state() survives torch.save / torch.load(weights_only=True); (5) Batch.validate() rejects an off-by-one label alignment; (6) validate() fires on each known-bad config in a table; (7) every Config field has a doc comment (A6 above).

2.4 The environment contract 

# align/envs/base.py  (sketch)
GenerateFn = Callable[[list[GenRequest]], list[Generation]]    # the rollout-engine seam; exists in lean_multiturn.py:152,160

class Env(Protocol):
    name: ClassVar[str]
    def specs(self, split: str) -> Sequence[EpisodeSpec]: ...
    def run_group(self, specs: Sequence[EpisodeSpec], generate: GenerateFn) -> list[EpisodeRecord]:
        """Run all episodes of one prompt-group in lock step so the engine can batch across them."""
    def rows(self, ep: EpisodeRecord) -> list[TrainingRow]:
        """response_mask is 1 only on tokens the generator returned; turn_ids labels each token."""
    def reward_components(self) -> tuple[str, ...]: ...
  • MathEnv (single turn): extract answer, OutcomeVerifier, CompositeReward; absorbs rl/assembly.py:57 assemble_rollout's scoring half.
  • LeanEnv(max_turns=K): rl/envs/lean_multiturn.py unchanged in behaviour; max_turns=1 is the single-turn Lean prover. Layout per_turn or concatenated stays an env option; PrefixBreak stays.
  • CodeEnv: new. Run candidate code against tests in a subprocess with the same constraints as verify/lean/sandbox.py (own session, prlimit, minimal environment, memory watchdog). Counts as a heavy job for the host preflight when pooled.

GrpoBatchSource becomes RolloutSource(env, engine, algorithm_assembler); MultiTurnBatchSource(GrpoBatchSource) (a subclass whose bind is "deliberately narrower than the parent's", rl/multiturn_source.py:418) disappears, because the parent no longer assumes problem/answer rows.

2.5 Exact mapping, current to new

F = in the import closure of the running P2 scripts (4.1): do not edit or move until the chain exits. For an F file the new path is created first as a facade that imports from the old path; the physical move and shim-direction flip happen in M7. Files not marked F can move in M3 to M5 with an identity shim.

Current New F
align/__init__.py stays; new lazy names in M7 F
align/config.py, exceptions.py, exit_codes.py, registry.py stay F
(new) align/protocols.py, align/batch.py
utils/dist.py, utils/guard.py, utils/tracking.py runtime/dist.py, runtime/guard.py, runtime/tracking.py F
utils/host_state_gate.py, utils/supervisor.py runtime/host_state_gate.py, runtime/supervisor.py
utils/__init__.py stays (shim package) F
kernels/* stay; new fused_policy_loss.py, paged_attention.py mixed
tokenizer/* stay F
model/{attention,checkpoint,hf_config,mistral,moe,moe_ep,olmoe,olmoe_checkpoint,qwen3,qwen3_checkpoint}.py stay F
model/{gates,weight_average}.py stay
model/tp.py [A2] stay F (imported by tests/test_tp_bench.py)
ppo/lora.py model/lora.py F
ppo/value_head.py model/heads.py F
train/qwen3_stage.py (build_qwen3_stage_model, resolve_qwen3, validate_qwen3_cfg, qwen3_trainer_extras, qwen3_entry_banner) model/build.py (builders) and recipes/sft.py (banner, extras) F
train/qwen3_stage.py LoRASFTModule algorithms/sft.py (LoRASFTModule) F
train/qwen3_stage.py Qwen3Format train/ckpt_format.py (next to DenseFormat, OlmoeFormat) F
train/{adamw8bit,quant8,offload_adamw}.py optim/{adamw8bit,quant8,offload_adamw}.py; make_optimizer to optim/factory.py F
train/{trainer,train_module,fsdp,probes,sharded_ckpt,ckpt_format,moe_module}.py stay (the engine); probes.py gains AdapterProbeView from ppo/trainer.py F
train/stage_gates.py, train/stage_smoke.py [A2] stay (they import only eval, which is L2b; moving them to recipes/ made dpo_stage import L5 and touched 12 frozen import sites for nothing). recipes/ imports them F
train/__init__.py, eval/__init__.py, model/__init__.py, rollout/__init__.py, kernels/__init__.py, verify/__init__.py, tokenizer/__init__.py, data/__init__.py, inference/__init__.py [A2] stay F
ppo/__init__.py, rl/__init__.py, rl/envs/__init__.py [A2: had no destination] stay as legacy alias packages. ppo/__init__ and rl/__init__ import 9 and 10 submodules eagerly (names in __all__), so after M7 each re-export line goes through a shim: they are rewritten as lazy __getattr__ re-exports to the new locations, in M7, with a test that from align.rl import GRPOTrainModule still works and that import align.rl imports no more than it does today F (rl/envs/__init__ free)
train/sft_loss.py algorithms/sft.py F
train/dpo_loss.py algorithms/dpo/loss.py
train/dpo_stage.py algorithms/dpo/{config,gates,source,adapter}.py; DpoBatch to batch.py
train/ref_cache.py algorithms/dpo/ref_cache.py
rl/types.py, TokenAdvRolloutBatch from rl/grpo_loss.py batch.py F
rl/grpo_loss.py algorithms/grpo/loss.py F
rl/grpo_module.py algorithms/grpo/module.py F
rl/advantages.py, rl/rewards.py, rl/ops.py functional/{advantages,rewards,ops}.py [A2: were algorithms/common/ and algorithms/grpo/, which put an L4 module under the L2 envs/lean] F
rl/tripwire.py, rl/canary.py algorithms/grpo/{tripwire,canary}.py F
rl/cispo.py algorithms/cispo.py (loss and conformance); fp32_head_* stay as the reference path for kernels/fused_policy_loss.py F
rl/stage_config.py algorithms/grpo/config.py
rl/batch_source.py, rl/assembly.py, rl/format_instrument.py algorithms/grpo/{source,assembly,format_instrument}.py
rl/multiturn_source.py algorithms/grpo/multiturn.py
rl/envs/lean_multiturn.py envs/lean/env.py; prompts vendored to envs/lean/prompts.py
ppo/gae.py functional/gae.py F
ppo/{loss,rewards,rollout,health}.py algorithms/ppo/{loss,rewards,rollout,health}.py F
ppo/trainer.py split: algorithms/ppo/module.py (PpoTrainModule), algorithms/ppo/source.py (PPOBatchSource), AdapterProbeView to train/probes.py F
ppo/stage_config.py algorithms/ppo/config.py
scripts/p2/expert_iteration.py library form algorithms/expert_iteration.py; script stays and later becomes a wrapper frozen
eval/verifier.py verify/outcome.py (old path becomes a shim). [A2] changes eval_source_hash, see 4.5 F
eval/{golden,harness,metrics,proof_stats,roundtrip}.py stay F
data/outcome_agreement.py eval/outcome_agreement.py ([A2] adds a file to align/eval/, so it changes eval_source_hash too: M7 only) F
data/* (rest) stay mixed
verify/{equiv,expr,extract,normalize}.py, verify/lean/* stay F (lean frozen by instruction)
rollout/protocol.py stays F
rollout/{engine,paged,tp_engine}.py [A2] stay; paged.py gets TritonPagedAttention in M6 F (imported by tests/test_tp_bench.py, which the chain launches)
rollout/{client,native}.py stay; new rollout/backends.py
inference/{engine,bucketed}.py, inference/receipts.py stay F / free
scripts/{sft,dpo,grpo,ppo}.py stay; recipes/* hold the same flows; scripts become 30-line wrappers after M7 F (sft)

3. Minimal public API

Ten call sites. Names marked exists are in the repo today; new are the facade layer from M1/M2. The point is that scripts/sft.py (261 lines), scripts/grpo.py (412) and scripts/ppo.py (482) are each one use of these ten.

import align

# 1. config (exists: align.config.load_config). The loader splits [algo], [env], [rollout] before struct-mode parsing, so
#    config.py itself is not edited during the freeze.
run = align.load_run("configs/p4_cispo_lean.toml", overrides=["optimizer.lr=1e-6"])        # new, wraps load_config

# 2. host guards, before torch (exists: align.utils.guard.apply_resource_guards)
align.runtime.apply_guards(cpu_threads=8, heavy=True, label="cispo-lean")

# 3. model + tokenizer, family chosen by run.model.arch through a registry (new; today scripts/sft.py::build_model
#    and train/qwen3_stage.py::build_qwen3_stage_model)
model, tok = align.model.build(run)

# 4. algorithm by name or preset (new; the classes wrap today's GRPOTrainModule / DpoStage / PpoTrainModule)
algo = align.algorithms.create(run.algo)            # run.algo = {name="grpo", preset="cispo_p4", eps_high=...}
refusals = algo.validate(run, world_size=3)         # launch refusals as data; recipes raise if non-empty

# 5. environment (new; wraps rl/envs/lean_multiturn.py)
env = align.envs.make("lean/multiturn", run.env, verifier=align.verify.lean.StrictVerifier(...))   # StrictVerifier exists

# 6. rollout engine (exists: align.rollout.client.make_rollout_client; new name and registry)
engine = align.rollout.make_engine(run.rollout, model=model, tokenizer=tok)          # "native" | "paged" | "tp" | "spool"

# 7. batch source: algorithm-specific, built from env + engine (new; GrpoBatchSource is today's equivalent)
source = algo.make_source(run.data, env=env, engine=engine, model=model, tokenizer=tok)

# 8. trainer: the existing Trainer, built from the algorithm (new classmethod over Trainer.__init__)
trainer = align.train.Trainer.from_algorithm(run, model, algo, source, run_dir=run.paths.run_dir)
trainer.resume_from_latest()                        # exists
summary = trainer.train()                           # exists; honours HALT codes and coordinated stop

# 9. one-liner for the common case (new)
summary = align.fit(run)                            # steps 2-8 in order, with the stage gates from train/stage_gates.py

# 10. kernels as a plain function, for people who only want the kernel (new)
loss, metrics = align.kernels.fused_policy_loss(
    hidden, lm_head_weight, labels, advantages, old_logprobs,
    objective="cispo", eps_high=4.0, beta=0.0, temperature=1.0, chunk_tokens=2048,
)

Evaluation is the eleventh and only other entry: align.eval.run(model, tasks=["minif2f"], n=32, verifier=...), which wraps eval/harness.py::run_eval.

CLI: python -m align grpo configs/p4_cispo_lean.toml --smoke calls recipes.grpo.run(...), so the existing torchrun command lines in the docs keep their shape.

What is deliberately not public: Trainer internals, TrainModule.optim_step, shard layouts, _-prefixed helpers. The lazy align/__init__.py pattern is kept (import align must not import torch; tests/test_public_api.py pins it).

4. Migration plan

4.1 The freeze closure

The task rule is: live GPU jobs run about two days; do not edit scripts/p2/*, extras/vllm_rollout/*, align/verify/lean/*, align/utils/guard.py, align/train/trainer.py. The P2 chain also launches new Python processes every round, and each one imports align.* fresh. A half-written or moved module breaks the next round even if the file itself is not on the list.

So I computed the static transitive import closure of the P2 entry points (scripts/p2/*.py, extras/vllm_rollout/*.py, scripts/sft.py, scripts/qwen3_export_check.py, 20 files), following every import including lazy ones and every parent package __init__ [read, AST script run for this document]. Result for those entry points alone: 101 of 129 align/ modules. It is a conservative over-approximation (a lazy import on a path a P2 round never executes still counts), which is the safe direction.

[A2] The first draft stopped there, but the chain scripts launch more than the P2 entry points. .tmp/p2/r1/chain2_gpu_validation.sh and chain_after_baseline.sh run tests/test_tp_bench.py, tests/test_moe_ep_bench.py, scripts/olmoe_sft_smoke.py and scripts/ckpt_convert.py as torchrun steps. Adding those four entry points pulls in four more modules (model/tp.py, rollout/engine.py, rollout/paged.py, rollout/tp_engine.py). Re-running the same AST closure gives 105 of 129 reachable, 24 free. This matters because M6 edits rollout/paged.py and M6/M7 both touch the TP engine. Rule: the entry-point list of the closure is every python/torchrun target named in any .tmp/p2/r1/chain*.sh that is still pending, and the closure test (test_frozen_surface.py) takes that list from a checked-in manifest, not from this document.

The 24 modules not in the closure, which may be edited or moved during the freeze:

align/data:      pretrain_corpus.py stageb_sources.py
align/inference: receipts.py
align/kernels:   _fa2_reference_upload.py fused_dpo.py ops.py
align/model:     gates.py weight_average.py
align/ppo:       stage_config.py
align/rl:        assembly.py batch_source.py format_instrument.py multiturn_source.py stage_config.py
align/rl/envs:   __init__.py lean_multiturn.py
align/rollout:   client.py native.py
align/train:     dpo_loss.py dpo_stage.py ref_cache.py
align/utils:     host_state_gate.py supervisor.py
align/verify/lean: sandbox.py

The names the frozen scripts and files import today (the contract the shims must keep) [read, AST]:

align.config: load_config                      align.registry: render_banner
align.data: hub; .ingest_parquet: iter_parquet_rows; .loader: collate_sft; .store: ShardDataset, ShardWriter
align.eval: proof_stats; .verifier: Outcome    align.exit_codes: PreflightRefused, run_with_exit_codes
align.model: qwen3, qwen3_checkpoint; .checkpoint: load_hf_safetensors
align.model.qwen3_checkpoint: align_to_hf, hf_to_align, read_qwen3_config, save_hf_qwen3
align.rollout.protocol: RolloutResult, read_request, spool_dirs, touch_heartbeat, write_result
align.train.qwen3_stage: build_qwen3_stage_model, qwen3_entry_banner, qwen3_trainer_extras
align.train.sft_loss: make_sft_loss_adapter    align.train.trainer: Trainer
align.train.stage_gates: GateLedger, assemble_sft_banner, assert_build_is_current, assert_distinct_devices,
    assert_probe_resolved, eval_source_hash, evaluate_and_gate, parse_args, record_golden_from_report,
    require_golden_receipt, run_entry_gates
align.train.stage_smoke: build_stage_model, build_training_sampler
align.utils: guard; .dist: barrier, init_distributed; .guard: apply_resource_guards, preflight_host_check
align.verify.lean: ReplPool, StrictConfig, StrictVerifier, find_toolchain, strict;
    .repl_pool: ReplError, ReplLimits, ReplProtocolError; .strict: AUDIT_PRELUDE, REF_NAME, lex, parse_statement

Names added by the four chain-launched entry points [A2]:

align.model: tp, moe, moe_ep, olmoe, olmoe_checkpoint; .qwen3: Qwen3Config
align.rollout.engine: PagedEngineConfig, SamplingParams     align.rollout.paged: kv_bytes_per_token, plan_capacity
align.rollout.tp_engine: TPPagedEngine
align.train: fsdp, sharded_ckpt; .ckpt_format: OlmoeFormat; .moe_module: build_moe_train_module; .offload_adamw: make_optimizer
align.train.trainer: MANDATORY_METRICS, Trainer, checkpoint_problem
align.utils.dist: local_rank, rank, seed_everything, world_size; .guard: assert_paths_off_root, host_state; .tracking: Tracker
align.data.store: ShardDataset      align.tokenizer: Tokenizer      align.model.olmoe_checkpoint: read_olmoe_config

tests/test_frozen_surface.py (M0) re-derives this list with the same AST scan on every run and asserts each name resolves, so a refactor that drops one fails in CI rather than in the middle of round 3. [A2] The scan's entry-point list is a checked-in manifest (tests/frozen_entrypoints.txt) that must contain every Python target named in a pending .tmp/p2/r1/chain*.sh.

4.2 Shim mechanics (verified)

For a file that moves, the old path becomes:

# align/rl/grpo_loss.py (after the move)
import sys
from align.algorithms.grpo import loss as _m
sys.modules[__name__] = _m

Checked with a throw-away package outside the repo [cpu-probe, deleted afterwards]: after import old.impl, old.impl is new.impl is True; from old.impl import _private_name, SomeClass returns the same objects; from old import impl returns the same module; unittest.mock.patch("old.impl.f", ...) changes the behaviour seen through the new path. So isinstance checks, private imports and monkeypatches in the distinct align.* import paths the tests use all keep working. (No test uses mock.patch("align..."), monkeypatch.setattr("align...") or import_module("align...") on a moved path today; 0 hits. [A2] re-checked by grep.) Pickle compatibility is not a concern: align checkpoints use safetensors and weights_only=True.

[A2] What the identity shim does NOT preserve (found by a second probe, same method).

  1. python -m old.path silently does nothing and exits 0. The runner imports the shim as __main__; the shim replaces sys.modules["__main__"] with the real module and the real module's if __name__ == "__main__": block never runs (its __name__ is the new dotted path). Probe: python -m pk.old.gate printed nothing and exited 0, python -m pk.new.gate printed REAL MAIN RAN and exited 3. align/utils/host_state_gate.py:31 documents python -m align.utils.host_state_gate and that module is the D-state safety gate (it must return EXIT_PREFLIGHT_REFUSED when the host is wedged); a silent exit 0 would turn the gate into a pass. Rule 3: a moved module that has an if __name__ == "__main__" block gets a shim that forwards it: if __name__ == "__main__": raise SystemExit(_m.main()). The manifest records which moved modules have a main, and test_legacy_shims.py runs python -m <old> in a subprocess for each and compares the exit code with python -m <new>. Today only host_state_gate.py among the moved modules has such a block (model/weight_average.py has a documented CLI but does not move).
  2. Tests that read a module's source text now read the shim. tests/test_dpo_loss.py:190 greps align/train/dpo_loss.py for from_pretrained, load_state_dict, AutoModel and NmoMistralForCausalLM; tests/test_h2_review_items.py:675-679 parses align/utils/host_state_gate.py for non-stdlib imports and reads align/utils/guard.py. After the shim they still pass, vacuously, because a 3-line shim contains none of those strings. See 4.6.
  3. Partial-initialisation cycles. ppo/__init__.py and rl/__init__.py import their submodules eagerly. An old-path shim executes from align.algorithms... import ..., which runs algorithms/__init__ and algorithms/<x>/__init__; if any of those import align.rl / align.ppo, the import order decides whether it works. The layering test cannot see this (shims are exempt). test_legacy_shims.py therefore imports every old and every new path in a fresh interpreter, new path first and old path first, and asserts both orders succeed.

Rules 1 and 2: shims warn only under ALIGN_WARN_LEGACY_IMPORTS=1 (a default warning would add stderr noise to every P2 round); and tests/test_legacy_shims.py asserts, from a manifest, that import_module(old) is import_module(new) for every moved path.

For an F file the direction is reversed first: the new module is a facade that imports from the old one. When the chain exits (M7) the real code moves and the old path becomes the shim above. The shim stays for at least one release because scripts/p2/* and extras/ import the old paths and are not ours to edit.

4.5 [A2] eval_source_hash: moving anything in or out of align/eval/ invalidates every gate ledger

train/stage_gates.py::eval_source_hash (:94-113) hashes the relative path and bytes of every .py under align/eval/. A GateLedger records that hash; on load, if the hash differs, _load deletes every gate in EVAL_DEPENDENT_GATES (the golden regression among them) and require_golden_receipt then refuses to train (stage_gates.py:739-753). Three planned moves change the hash: eval/verifier.py becoming a shim, data/outcome_agreement.py arriving in eval/, and the M8 path-hygiene edit of eval/golden.py. The first draft treated these as free once the chain had exited. They are not: every run directory that is resumed after M7 needs the golden regression re-run (a 7B checkpoint job; scripts/eval.py --mode=golden).

Plan: (a) none of the three edits happens before the P2 chain exits (already true for M7, now stated for M8); (b) M7 and M8 land as one hash epoch, i.e. adjacent commits with a single golden re-run after the second; (c) tests/test_layering.py and test_legacy_shims.py print the new eval_source_hash so the epoch boundary is visible in CI; (d) the alternative that avoids the epoch keeps align/eval/ byte-identical: leave Outcome/OutcomeVerifier in eval/verifier.py; cut verify -> eval (verify/lean/reward.py:116) by typing the result against a small Protocol in verify/; and put outcome_agreement.py in recipes/ (L5, which may import eval) instead of eval/, with align.data.__getattr__ serving the old name lazily. It costs a slightly less tidy layout and saves a 7B golden re-run, so it is the cheaper choice unless the owner wants the tidier layout. The owner decides (open question 11).

4.6 [A2] Tests that pin source text at old paths

Reading file text of a library module or of scripts/*.py is an established pattern here (the CPU suite cannot execute the GPU scripts, so it pins them by structure). Found by grep: test_dpo_resume_wiring.py (string-matches scripts/dpo.py: '"--resume"', os.environ.get("ALIGN_RESUME") == "1", the rotation guard), test_grpo_resume.py:293 (AST of scripts/grpo.py: resume_from_latest call ordering), test_dpo_stage.py:1566 (assert_launch_topology(cfg) before run_entry_gates(), test_ref_cache.py:179,300 (scripts/dpo.py, scripts/precompute_ref_logprobs.py), test_h2_review_items.py:374,383 (train() in scripts/sft.py, pretrain.py), test_qwen3_stage.py:301, test_g0_integration.py:38, test_olmoe_scripts.py (several), test_verifier.py:371 (scripts/eval.py), plus the library-text pins in 4.2 item 2.

Consequences for the plan: (1) M7 "make scripts/* 30-line wrappers" cannot keep the baseline green unless each pin is retargeted to the recipes/ file that now holds the flow, with the same assertion, in the same commit. M2's recipes are copies, so the pins should be written against both files from M2 on (a parametrised [script, recipe] pin), and the script pin is dropped only at M7. (2) M3's move of train/dpo_loss.py (and utils/host_state_gate.py) must retarget the two library-text pins to the new path in the same commit, and the shim manifest records the pin so test_legacy_shims.py fails if a pinned path is a shim (a vacuous pin is a lost guard, not a pass). (3) The gate "same passed counts" cannot hold literally once contract tests are added; it is restated in 4.3 as "every baseline test id still passes".

4.3 Steps

Every step has the same gate: pytest -m "not gpu and not oracle and not slow" where every test id in the M0 junit baseline still passes or is skipped for the same reason ([A2] not "equal counts": M1 onward adds contract tests, so the totals must grow, while a baseline id going missing or silently skipping must fail), tests/test_dependency_firewall.py green, tests/test_public_api.py green (import stays cheap), tests/test_frozen_surface.py green, and, from M3, tests/test_legacy_shims.py (identity, -m forwarding, both import orders, no vacuous source pin: 4.2 and 4.6). All GPU-free. All run with CUDA_VISIBLE_DEVICES=. [A2] The suite shares a 48-core box with the live chain: run it under nice -n 10, OMP_NUM_THREADS=4, -p no:cacheprovider, and only after the host preflight in /apps/CLAUDE.md (free -g; uptime) passes; conftest.py already keeps tmp_path off /tmp.

Step What Effort Touches F files? Extra gate
M0 Safety net, no code moves: record the real CPU baseline (junit file); add test_frozen_surface.py, test_legacy_shims.py (empty manifest), test_layering.py (report-only, lists today's violations from 2.1) 0.5 d no baseline numbers committed in the doc
M1 Facade layer, additive: protocols.py, batch.py (re-exporting the existing batch classes), algorithms/{base,registry}.py and algorithms/{sft,dpo,ppo,grpo,cispo,dapo}.py as classes that wrap the existing loss adapters and TrainModule subclasses; contract test suite 2.3 1.5 d no (imports only) contract tests pass for all six
M2 recipes/ and cli.py: the main() flows of scripts/{sft,dpo,grpo,ppo}.py as functions, written as new files (the scripts are not edited). align.load_run splits [algo]/[env]/[rollout] before load_config. 3 d no tiny-model parity: recipes.sft and scripts/sft.py --smoke produce identical history.jsonl losses on CPU/gloo (train/stage_smoke.py already builds the model); same for grpo and dpo smoke
M3 Move the non-F modules with identity shims: dpo_*, ref_cache, rl/{stage_config,batch_source,assembly,format_instrument,multiturn_source}, ppo/stage_config, utils/{host_state_gate,supervisor}. Fix cycles that need only free files. [A2] Retarget the two library-text pins (test_dpo_loss.py:190, test_h2_review_items.py:675) in the same commit; host_state_gate gets the __main__-forwarding shim 2 d no test_legacy_shims.py manifest grows by exactly the moved paths; python -m align.utils.host_state_gate exit code equals the pre-move exit code
M4 envs/: Env protocol, MathEnv, LeanEnv from rl/envs/lean_multiturn.py; vendor lean_prompts.py as envs/lean/prompts.py with a sha256 identity test against extras/vllm_rollout/lean_prompts.py; extract displayed_statement_matches; collapse MultiTurnBatchSource into RolloutSource(env=...) 3 d no tests/test_lean_multiturn_env.py, test_multiturn_source.py, test_lean_gen_contract.py unchanged and green; GRPO smoke bit-equal before and after
M5 Kernel A/B: kernels/fused_policy_loss.py as a chunked autograd.Function (5.4-A) with an eager fp64 oracle test; mm(out_dtype=float32) probe on GPU after the freeze 3 d CPU + 0.5 d GPU no fp64 gradient oracle (max abs error recorded); GPU parity vs fp32_head_logprobs
M6 Kernel C: Triton paged-attention decode (5.4-C) behind TritonPagedAttention, fast and invariant modes 8 d (CPU interpreter tests 3 d; GPU validation and tuning 5 d) [A2] yes: rollout/paged.py is in the closure of tests/test_tp_bench.py; starts only after the chain exits bit-invariance test: row output independent of batch composition; tests/test_paged_engine*.py green; decode benchmark vs TorchPagedAttention
M7 After the P2 chain has exited (gate [A2]: pgrep -f '\.tmp/p2/r1/chain' and pgrep -f 'scripts/p2/' both empty, not only chain_after_baseline: chain4.sh and chain5.sh are live today; nvidia-smi idle, as in docs/P3_ENGINE.md GPU plan): move the F files, flip each facade to a shim, rewrite the three legacy __init__s as lazy re-exports, add the new lazy names to align/__init__.py, retarget the script-text pins (4.6) and make scripts/* 30-line wrappers, sweep path mentions in docs/ (README, ARCHITECTURE.md, DEPENDENCY_POLICY.md and 6 more name align/rl/, align/ppo/, align/utils/, align/train/...) 3 d (was 1.5: about 55 F modules, ten-odd pins, three __init__s, docs) yes full suite and frozen-surface test after each moved file, not at the end; then one golden re-run (4.5)
M8 Open-source hygiene: remove the hard-coded paths (2.1), Nmo* names (aliases kept; 38 files mention nmo/NMO), accept both ALIGN_* and NMO_* env vars, add a LICENSE (none exists in the repo root), relax torch==2.11.0 / triton==3.6.0 / requires-python <3.13 pins to tested ranges, CI lane with ruff + firewall + layering, README rewritten around the ten call sites 3 d yes clean checkout test: uv venv && uv pip install -e .[dev] and the CPU suite pass with no /apps path present

Total about 27.5 working days (about 5.5 weeks for one engineer; [A2] M7 re-estimated from 1.5 to 3 d). M1 to M5 are all runnable during the freeze because they only add files or move non-F ones. M6 (it edits the frozen rollout/paged.py) and the GPU halves of M5 wait for free GPUs.

4.4 Risks

Risk Mitigation
A new file in an existing package changes import order or a __init__ side effect Only new subpackages are added during the freeze; no existing __init__.py is edited until M7
The facade and the old class drift while both exist Facades hold no logic; contract tests call both paths and compare
sys.modules shim hides a circular import test_layering.py is strict before M7 flips anything
M2 recipes diverge from scripts that P2 still runs Parity test on tiny model; the scripts are the source of truth until M7; the source pins of 4.6 are parametrised over [script, recipe] from M2
Test counts in docs are stale, so "stay green" has no baseline M0 records it
Moving train/stage_gates.py [A2] No longer moved (2.1). eval/ changes still invalidate the ledger, see 4.5
[A2] python -m <old.path> through a sys.modules shim silently does nothing and exits 0 host_state_gate.py is a safety gate with a documented python -m align.utils.host_state_gate. Shim rule 3 in 4.2 plus a subprocess test
[A2] Source-pinning tests read old file paths and script text (4.6) Retarget each pin in the same commit as the move, and fail the gate if a pin goes vacuous
[A2] Eval source hash epoch (4.5) Plan one golden re-run in M7

5. Where custom Triton kernels pay off

5.1 Method and constants

For each candidate the question is: is the op bound by bytes or by FLOPs, and how much of the iteration is it? A kernel can only recover the gap between what the code does today and the roofline. Constants:

Constant Value Source
A10G HBM bandwidth 600 GB/s [assumed] datasheet, same value used in docs/P3_ENGINE.md
bf16 tensor-core peak / best measured about 70 / 68 TFLOP/s (one 8192x4096 @ 4096x4096); 44 to 53 TFLOP/s in-layer [read] docs/MEASUREMENTS.md:102-105,254
fp32 CUDA-core peak, sustained 31.2 peak, 25 sustained (sensitivity 20 to 31) [assumed] datasheet; not measured here. TF32 is disabled in rl/cispo.py::_strict_fp32_matmul
Qwen3-8B (Goedel-Prover-V2-8B) hidden 4096, 36 layers, 32 q heads / 8 kv heads, head_dim 128, vocab 151,936, intermediate 12,288, untied head [read] model config.json
KV bytes per token 2 x 36 x 8 x 128 x 2 B = 147,456 B (144 KiB) [calc]; matches docs/P3_ENGINE.md (3.60 GiB / 26,192 tokens)
Weights bf16 15.26 GiB = 16.39 GB; 27.3 ms to read once at 600 GB/s [calc]; 15.26 GiB is docs/P3_ENGINE.md
Eager launch cost 6 us per op [assumed], typical 5 to 10 us

5.2 Where the time goes in one RL iteration (why decode dominates)

Illustration with stated knobs, not a measurement. 16 prompts x 8 samples, mean completion 4,096 tokens (cap 16,384): 524k generated tokens per iteration.

  • Generation at vLLM-class throughput (about 876 tok/s end to end, the vLLM figure quoted in docs/P3_ENGINE.md): about 600 s. At the pure-torch engine's modelled tp=4 rate (218 to 386 tok/s, rows 3 to 5 of the table in 5.4-C; [A2] the first draft's low end, 100 tok/s, came from a tp=1 batch of 16 that does not fit in memory, see 5.4-C): 1,360 to 2,400 s.
  • Training on the same tokens, LoRA with activation checkpointing: body forward 2 x 6.95e9 x tokens (the 36 blocks hold 6.95 B parameters [calc]); with recompute and a dX-only backward, about 3x forward = 41.7 GFLOP/token; 524k tokens = 21.8 PFLOP; on 3 train GPUs at 45 TFLOP/s with no comm: about 160 s. FSDP over host shared memory is comm-bound on this box (docs/MEASUREMENTS.md section (h), 92% of a micro-step), so 160 s is a floor.

Reading: with the native engine, generation is about 8x to 15x the training body compute (the 160 s above excludes the head, which adds about 35 s on 3 GPUs: 32 micro-batches x 3.3 s / 3, see 5.4-A); even against vLLM it is the larger term. The decode path is where a kernel changes the iteration time; the head is a second-order term. This ordering drives 5.5.

5.3 What exists

  • Triton: kernels/flash_attention.py (FA2 forward and backward, adapted from the user's Flash-Attention-101/triton/flash_attention.py, kept as _fa2_reference_upload.py), kernels/ops.py (RMSNorm/RoPE/SwiGLU, pedagogy, off by default). The FA2 backend is not selectable at HEAD: no gate artifact at the gate path (REFERENCE_GAP_ANALYSIS 1.1 item 9); SDPA is the working backend.
  • Not Triton, despite the name: kernels/fused_ce.py and kernels/fused_dpo.py are chunked torch.autograd.Functions. They implement the Liger FusedLinearCrossEntropy design (arXiv:2410.10989, "FusedLinearCrossEntropy" paragraph, Eq. 15-17: chunk the head, write dh = W dlogits and dW += h dlogits^T during forward) in torch matmuls.
  • rollout/paged.py::TritonPagedAttention is an interface with raise NotImplementedError. Its docstring already states the kernel contract, including bit-invariance ("fixed KV-split count per context length, no atomics").

5.4 Candidates

A. Fused head + policy-gradient token loss (with logprob gather)

Today (rl/cispo.py::fp32_head_logprobs, :202): per chunk of 2,048 tokens, h.float() @ W32.T in strict fp32, logsumexp, gather; torch.utils.checkpoint per chunk, so backward recomputes the GEMM; a separate no-grad pass for entropy.

GEMM count at 2*H*V = 1.245 GFLOP/token (20.4 TFLOP per 16,384-token micro-batch [calc]), frozen head (LoRA):

Path GEMMs Arithmetic Time per 16k tokens
Today: fwd + checkpoint recompute + dh + entropy pass 4, strict fp32 4 x 20.4 / 25 3.3 s (2.6 to 4.1 for 31 to 20 TFLOP/s)
Fused: fwd with dlogits formed in-chunk, then dh = dlogits @ W; entropy from the same sweep 2, bf16 operands, fp32 output 2 x 20.4 / 50 0.82 s (0.68 to 0.91 for 60 to 45)

Ratio about 4x [calc, with the assumed fp32 rate]. If the head is trainable (full FT), add one dW GEMM to each side (5 vs 3): the ratio is the same.

Why the in-chunk form works for this loss. For CISPO the per-token gradient of the objective J with respect to the logit vector is c_t (softmax(z) - onehot(y)) with c_t = -A_t * sg(clip(r_t, ., 1+eps_high)) (CISPO Eq. 4-5, arXiv:2506.13585 Sec. 3.1; the weight is a stop-gradient constant). [A2] The sign is that of the ascent direction: the loss -J that the optimizer minimises has c_t = +A_t * sg(clip(r_t, ., 1+eps_high)). Checked in fp64 autograd on a 11-way softmax: both signs reproduce the autograd gradient to 7e-18, and the two differ only by the overall sign. The kernel must use the loss sign. r_t needs the new logprob, which the same chunk has just computed, so c_t is known before dlogits is formed. The KL term adds a closed-form function of the same logprob. This is exactly the structure of fused_ce._FusedLinearCrossEntropy (gradient written in forward, backward is a scalar rescale), extended with the token coefficient c_t. For the clipped-surrogate objectives (DAPO/A3) c_t = -A_t r_t 1[not clipped], also known in-chunk.

The precision insight (the real reason the fp32 head matters). MiniMax (arXiv:2506.13585 Sec. 3.2) attribute the train/inference probability mismatch to "high-magnitude activations in the LM head" and fix it with an FP32 head (correlation about 0.987 to 0.997, their Fig. 3). In align the head weights and the final hidden state are bf16. A bf16 product is exact in fp32 (8+8 mantissa bits < 24), and both a bf16 tensor-core GEMM and an fp32 GEMM accumulate in fp32. So an "fp32 head" built from bf16 operands differs from a bf16 GEMM only in whether the output logits are rounded to bf16. At |z| ~ 16..32 the bf16 spacing is 0.125, so a logit is off by up to 0.06 and so is a logprob.

[cpu-probe] N=512, H=1024, V=16,384, random bf16 h and W with logits up to |z| = 50.6; fp64 reference of the same bf16 operands. fp32-head max abs logprob error 7.1e-6; bf16-output matmul max abs logprob error 0.145 (mean 0.049). That is the whole effect, reproduced on CPU, with operands identical in both arms.

Consequence: the head does not need fp32 compute; it needs fp32 output. torch.mm(a_bf16, b_bf16, out_dtype=torch.float32) exists in torch 2.11 (the op aten::mm.dtype is registered; on CPU it reports "could not run on the CPU backend", so it is CUDA-only and unverified on a GPU here). The current fused_ce.py: line logits = torch.matmul(h_chunk, weight.t()).to(cdtype) rounds to bf16 first and upcasts after, so the SFT loss carries this rounding noise too (harmless for SFT gradients, which are unbiased in aggregate; it matters for ratios). The same trick makes the P4 fp32 head about 3x to 4x cheaper than strict-fp32 and also saves the fp32 weight copy (weight.float() is 2.49 GB, against 1.25 GB bf16 [calc]) on a 23 GB card.

Caveat to test on GPU: tensor-core accumulation is not IEEE fp32 (reduced alignment and truncation inside the MMA). The size of that error against bf16 output rounding must be measured with the probe above run on CUDA.

[A2] The backward GEMM is not "free precision" either. The second fused GEMM, dh = dlogits @ W, only runs at the bf16 tensor-core rate if dlogits is cast to bf16 first (a strict-fp32 dlogits @ W is exactly the cost being removed). The forward logprob keeps the 7e-6 accuracy above, but dh does not. [cpu-probe] N=512, H=1024, V=16,384, random bf16 h, W, token coefficients c_t ~ N(0,1): relative L2 error of dh against fp64 is 3e-7 with an fp32 dlogits and 1.8e-3 with a bf16-rounded dlogits (max abs 1.8e-3 against max|dh| 0.68). That is ordinary bf16 backward noise, but it means the fused path is not bit-comparable with the strict-fp32 path, the M5 oracle test needs a tolerance of about 5e-3 relative on dh (and 1e-4 on logprobs), and "fp32 head" in the P4 pre-registration should be read as "fp32 logits", not "fp32 gradient". The cast adds one read and one write of the tile (about 1.9 GB per 2,048-row chunk, 25 ms per 16k tokens), already inside the 10% to 16% range below.

Share of a training step. Body about 3 x 228 TFLOP at 45 TFLOP/s is about 15 s per 16,384-token micro-batch [calc]; head 3.3 s today versus 0.82 s fused: a saving of 2.5 s of about 18.3 s, about 13% of the training half of an RL iteration, and about [A2] 1% to 3% of the iteration end to end (32 micro-batches x 2.5 s / 3 GPUs = 27 s against an iteration of about 1,550 to 2,600 s with the native engine, about 800 s against vLLM; 5.2). It is worth doing because it is cheap and also removes 1.2 GB of resident memory, not because it moves the iteration.

Is Triton needed? No. Elementwise traffic per 2,048-row chunk: fp32 logits are 1.245 GB; today about 8 passes (GEMM write, logsumexp read, recompute write and read, softmax write and read, dlogits read) is 10 GB, or 16.6 ms at 600 GB/s, 133 ms per 16k tokens [calc]. Against today's 3.3 s that is 4%. [A2] The denominator that matters for the Triton decision is the fused path's, not today's. After A the head costs 0.82 s, and its elementwise tail is about 7 passes (logits write, logsumexp read, softmax read, dlogits write, bf16 cast read and write, dh GEMM read, about 7.5 GB per chunk, 100 ms per 16k tokens): 10% to 16% of the fused head, 0.5% to 0.7% of the 18.3 s training step. A 3x faster fused epilogue (the Liger figure below) would recover about two thirds of that, at most 0.07 s per micro-batch. The conclusion (do not write Triton for it) stands, but at the stated 4% the reasoning understated the residual by 3x. Fusing it (Liger reports a 3x faster cross-entropy kernel in isolation, arXiv:2410.10989 Fig. 2a and 3a, at vocab 163,840, benchmarking the kernel alone with the GEMM excluded) would recover at most about 0.1 s per micro-batch. AgileRL's own note that its Liger GRPO kernel gives no speedup (1.2) is the same fact. Ship it as a chunked autograd.Function (kernels/fused_policy_loss.py), tested against fp32_head_logprobs and an fp64 oracle.

B. fp32-head logsumexp as a standalone kernel

Does not pay. After A it is an epilogue over a [chunk, V] fp32 tile, about 10% to 16% of the fused head's time (4% of today's), and a Triton GEMM on sm_86 (99 KB shared memory, docs/MEASUREMENTS.md:124) is unlikely to beat cuBLAS on these shapes. The same document reached the same conclusion for the grouped-GEMM kernel (P1-12 fails its own entry condition). The only reason to write it later is to avoid materialising the [chunk, V] tile at all (memory, not time), and chunk size already bounds that: at 512 rows the tile is 311 MB.

C. Paged-attention decode (the one that pays)

Today (rollout/paged.py::TorchPagedAttention, 'segment' mode, read above): for every decode segment, every layer: cache.gather copies K and V into a contiguous strip (read L + write L), a second copy into a zero-filled padded buffer (zero fill L_max, copy read L + write L), then one SDPA over L_max with a boolean mask. Per sequence per layer per tensor (K or V) the traffic is about 4L + 2 L_max, against L for a fused kernel that reads blocks once through the block table. With rho = L_max / L_mean, the multiplier is 4 + 2 rho: 8x at rho = 2. The length distribution of a prover RL batch is heavy-tailed (most stop in a few thousand tokens, a few reach the 16k to 40k cap, rollout/paged.py docstring), so rho >= 2 is the normal case, and docs/P3_ENGINE.md already flags the padded transient (1.17 GiB per layer call at 64 sequences, 4.7 GiB if SDPA falls back to the math backend with GQA expansion).

On top of bytes there is launch overhead: about 12 eager ops per segment per layer is B x 36 x 12 launches per step (27,648 at B = 64), about 166 ms at 6 us [assumed]; the GPU is idle while Python issues them.

Model (all [calc] from 5.1; tp as in docs/P3_ENGINE.md; comm term is that doc's assumed 20 to 26 ms for B = 64 at tp = 4 and about 10 ms for B = 16; fused kernel at 85% of HBM peak [assumed]; step = max(GPU time, CPU launch time) for today):

Setting KV tokens (capacity) Weights ms KV ideal ms Today step ms (tok/s) Fused step ms (tok/s) Speedup
tp=1, B=6, ctx 4,096, rho=2 24,576 (26,192) 27.3 6.0 76 (79) 34 (174) 2.2x
tp=4, B=16, ctx 4,096, rho=2 65,536 (438,048) 6.8 4.0 49 (326) 22 (743) 2.3x
tp=4, B=64, ctx 4,096, rho=2 262,144 (438,048) 6.8 16.1 166 (386) 52 (1,240) 3.2x
tp=4, B=48, ctx 8,192, rho=2 393,216 (438,048) 6.8 24.2 220 (218) 55 (870) 4.0x

[A2] Rows removed from the first draft because they do not fit in the card. The first draft listed tp=1 B=16 at ctx 4,096 and 8,192 (3.4x and 4.4x) and tp=4 B=64 at ctx 8,192 (4.2x). The KV pool the native engine can allocate is in docs/P3_ENGINE.md ("Native engine, bf16 KV" table): 26,192 tokens at tp=1 (weights 15.26 GiB leave 3.60 GiB), 163,472 at tp=2, 438,048 at tp=4 (util 0.90). tp=1 B=16 at ctx 4,096 needs 65,536 tokens (2.5x over) and at 8,192 needs 131,072 (5x over); tp=4 B=64 at 8,192 needs 524,288 (1.2x over). Those speedups were computed correctly for the stated inputs but describe a batch the engine cannot hold, and the 3.4x to 4.4x figures were the top of the quoted range. Replacement rows use the largest batch that fits (B=6 at tp=1, B=48 at tp=4 with ctx 8,192; 20 ms comm for B=48). Row arithmetic, checked: KV ideal = tokens x 147,456 B / tp / 600 GB/s; today = weights

  • (4+2rho) x KV ideal + comm (launch-bound at B=64: 27,648 x 6 us = 166 ms); fused = weights + KV ideal / 0.85 + comm.

Expected: 2x to 4x on decode steps at the largest batch that fits (6 to 64) and rho about 2; larger with a larger rho or longer context; at batch 1 the step is weight-bound and the kernel changes little (it still removes the launch overhead, 2.6 ms vs 27 ms). Amdahl: this speeds the generation half, which 5.2 shows is the larger term. The bar in docs/P3_ENGINE.md (native at least 0.5x vLLM throughput) is the acceptance test, and the stated fallback there ("the Triton kernel is the next step, not TP") is consistent with this analysis. The vLLM kernels in the PagedAttention paper (arXiv:2309.06180 Sec. 5.1, "Kernel-level Optimization") are three fusions: fused reshape-and-block-write ("to minimize kernel launch overheads"), fused block read and attention (adapted from the FasterTransformer attention kernel, reading the KV cache through the block table), and fused block copy.

Design constraint from RL, not performance: the per-row key reduction must not depend on batch composition (the invariant mode of TorchPagedAttention, rollout/paged.py). Use a split-KV (flash-decoding style) kernel with a fixed split size per context length (e.g. 512 keys per split, so the split count is ceil(L / 512), a function of L only), a second reduce kernel in a fixed order, no atomics. That keeps rho (the train/infer logprob ratio) unaffected by who shares the batch. It is the user's FA2 forward loop (Flash-Attention-101/triton/flash_attention.py, online softmax per KV block) with Lq = 1, the KV loop split across programs, and the block table as an extra indirection; the numbered-note house style applies (one note on split-KV online softmax merging, one on block-table indexing, each with a notebook that checks the kernel against the naive version).

Validation plan (needs free GPUs): (1) bitwise: same sequence, run alone and inside batches of 2, 16, 64 with different neighbours: identical logits. (2) numeric: against TorchPagedAttention(mode="row") at bf16, max abs logit error stated. (3) throughput: tests/test_paged_engine_bench.py --gpu with the same prompt ids as the vLLM manifest. (4) peak memory with --reserve-gib reduced, since the padded transient goes away.

D. Candidates that do not pay

Candidate Estimate Verdict
Per-token CISPO/GRPO loss arithmetic alone tensors of B x T fp32: 64 KB at 16k tokens; about 10 elementwise ops is below 1 MB of traffic, about 50 us of launch time against a multi-second step Do not write a kernel
adamw8bit Triton step (open item in DEPENDENCY_POLICY) ideal about 18 GB of traffic for 1.8 B local elements [calc, about 10 B/element] = 30 ms; runs once per optimizer step, and a step is 8 to 85 micro-steps of seconds each Under 1% of a step. No
Fused LoRA add (W x + B A x) one saved read+write of each [T, out] bf16 activation per projection: 2.55 GB per layer at 16k tokens = 4.3 ms, 153 ms per forward over 36 layers [calc] About 1% of a 15 s body pass; no
FA2 Triton as default Gate exists and is not enabled (no artifact at the gate path); attention is about 4% of layer FLOPs at seq 1k (kernels/fused_ce.py docstring), more at 16k Finish the equivalence gate first; a measurement task, not a kernel task
RMSNorm/RoPE/SwiGLU Triton (kernels/ops.py) kept flag-gated and off; torch.compile fuses them Keep as pedagogy

5.5 Ranking

  1. C, paged-attention decode. About 2x to 4x on decode (3x central), on the largest term. 8 days including GPU tuning. Needs free GPUs.
  2. A, fused head + policy loss as an autograd.Function with out_dtype=float32. About 4x on the head, about 13% of the training half, 1.2 GB less resident memory, and it fixes the SFT output-rounding as a side effect. 3 days plus 0.5 day GPU. No Triton.
  3. Everything else in 5.4-D: no.

6. Open questions for the owner

  1. Name. align is a common PyPI name. An open-source release needs a distinct distribution name; the import name can stay. Renaming Nmo* (38 files) is mechanical; do you want the old names kept as aliases for one release?
  2. Licence. The repo has no LICENSE. AgileRL is Apache-2.0, Liger-Kernel is BSD-2-Clause, open-instruct (vendored GRPO loss, rl/grpo_loss.py) is Apache-2.0: Apache-2.0 for align is compatible. Confirm. [A2] The vendored open-instruct code (rl/__init__.py says "vendored verbatim") also needs its upstream copyright notices kept in the files and a NOTICE entry before release; M8 does not list that, add it.
  3. Is expert_iteration an algorithm or a recipe for you? The proposal makes it an OuterLoop with its own small interface (2.3). If you would rather keep it a script, drop it from the registry and keep only recipes.sft and recipes.eval.
  4. A10G constants. Section 5 rests on 600 GB/s and a 25 TFLOP/s sustained fp32 rate, both assumed, plus 6 us eager launch cost. Three 2-minute microbenchmarks (a bandwidth copy, an fp32 and a bf16 GEMM at the head shape, an empty-op loop) on a free GPU would turn every [assumed] in 5.1 into a measurement. They are GPU jobs and must take the heavy-job lock.
  5. mm(out_dtype=float32) on CUDA. Not runnable here. If it is absent or slow for the head shape the fallback is a Triton GEMM with fp32 epilogue, which moves A from 3 days to about 6 and changes the 5.5 ranking only slightly.
  6. Tensor-core accumulation error versus bf16 output rounding for the head (5.4-A caveat): needs the CPU probe repeated on a GPU.
  7. Test layout. AgileRL mirrors the package in tests/. This proposal keeps the 121 flat files and adds tests/contract/. Mirroring would change documented pytest tests/test_x.py commands in several docs; say if you want it.
  8. Which [algo] name for GSPO-style trajectory ratios. Not implemented in align. Add it in M1 or leave it out until a run needs it?
  9. inference/engine.py (contiguous KV) versus PagedEngine. After M6 the paged engine can serve eval too, and inference/ could be deleted (962 + 202 + 220 lines). Not proposed now because eval receipts and golden regression depend on its exact numerics.
  10. Test baseline. README says 1608 passed, DEPENDENCY_POLICY 1089, REFERENCE_GAP_ANALYSIS 1107. M0 resolves it; until then no "stays green" claim here has a number behind it.
  11. [A2] Eval hash epoch. Moving Outcome into verify/ and outcome_agreement.py into eval/, and the M8 edit of eval/golden.py, change eval_source_hash and invalidate every gate ledger, so each run dir resumed afterwards needs a golden re-run (4.5). Accept one re-run at M7/M8, or take the alternative that keeps align/eval/ byte-identical?
  12. [A2] Chain manifest. The freeze closure is only as complete as the list of entry points. Should the repo carry tests/frozen_entrypoints.txt, maintained by whoever writes a new .tmp/p2/r1/chain*.sh, or should the closure scan parse the chain scripts directly?
  13. [A2] CodeEnv binaries. prlimit and whatever interpreter runs candidate code are external binaries; DEPENDENCY_POLICY requires a dated entry for those (as for Lean, 2026-10-12). Which languages should CodeEnv support first?