EmergentTwilight/MMCoT/OURS - Storage Bucket

UniAR backbone, interleaved SFT, OPD, and GRPO

Third-batch GRPO configs live in config/rl/RLv3/. The A100×4 baseline uses dynamic rollout/validation queues and checkpointed SIQ one-hour attempts. Policy replay checkpoints vocabulary/JSD projections in 128-token chunks; parent-linked execution migrations preserve optimizer/reference state while repairing memory use on long trajectories. Validation commits full hard trajectories and scores per query, so retries reuse finished samples. See UniAR runtime contracts, submission snapshots for recovery, and modeling for the SFT/OPD/RL method.

OURS is the project-owned final library. It contains the independent UniAR backbone plus reusable shared/SFT/OPD/GRPO code, but no generated outputs, viewers, external repositories, or imports from projects/UniAR/experiments.

Install

From the repository root:

mamba run -n UniAR python -m pip install -e 'OURS[train,decoder]'

The released weights remain data and can be loaded from models/UniAR-SFT with from_pretrained().

Custom SFT losses

forward_interleaved() teacher-forces mixed text/BSQ trajectories without choosing a training objective. It returns pre-final-norm language states, Qwen3-VL M-RoPE positions, the visual-position mask, and visual-head logits:

output = model.forward_interleaved(
    input_ids=input_ids,
    attention_mask=attention_mask,
    image_grid_thw=image_grid_thw,
    visual_codes=visual_codes,
)
text_logits = model.project_text(output.hidden_states)
loss = custom_loss(text_logits, output.visual_logits, output)
loss.backward()

The package preserves the released checkpoint names and UniAR serialization: canonical ChatML, <image_gen> generate image H W<|vision_start|>, grouped BSQ visual positions, <|vision_end|>, and the fast-tokenizer special-token IDs.

SFT framework

The canonical configuration is config/sft/zebra_ig_lora.yaml. It contains every model, data, LoRA, text CE, bit CE, optional Soft-BSQ features, bidirectional Information Gain, EMA Necessity distillation, optimizer, scheduler, evaluation, checkpoint, and logging setting. Loading is strict: unknown keys, wrong primitive/Literal types, and incompatible IG settings fail before model loading. This first framework is intentionally LoRA/rsLoRA-only; it rejects a misleading unsharded BF16 full-parameter mode.

The current v1 Tetris pair is config/sft/v1/zebra_tetris_{vanilla,ig}_base.yaml. Vanilla uses the canonical half-sample partition with trajectory/update-uniform BitFlip up to 0.1, while its controlled IG counterpart changes only the IG and validation switches. Relative IG uses all events, unit weights, and temperature 0.25 in both directions. Derived searches replace base in both filename and run.name with the changed variable or schedule.

The first Visual-Jigsaw v1 comparison uses zebra_visual_jigsaw_{vanilla,ig}_3epoch.yaml. It keeps the same recipe on the full Visual-Jigsaw split; the IG file is the initial $\lambda=1$ point for later weight searches.

config/sft/v1/thinkmorph_vanilla_base.yaml applies the same UniAR vanilla recipe to all 24,990 ThinkMorph training rows in data/ThinkMorph/ThinkMorph. It disables internal validation and declares no validation/test split; see docs/thinkmorph.md for the corpus and the separate BAGEL baseline.

config/sft/v2/ is the constant-LR family. Its migrated base configs use dataset-fixed warmup updates rather than a fraction of the final training horizon. thinkmorph_vanilla_r64_val5_bitlambda3_16epoch.yaml runs 16 epochs on IDVal5, saves every exact epoch boundary, and validates only at those boundaries, enabling representative 4/8/12/16-epoch evaluation from one training trajectory.

PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m sft inspect \
  --config OURS/config/sft/zebra_ig_lora.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m sft probe \
  --config OURS/config/sft/zebra_ig_lora.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR torchrun --nproc-per-node=4 \
  -m sft train --config OURS/config/sft/zebra_ig_lora.yaml

inspect is CPU-only. probe reports full/blind/gain scores for three proven members of the materialized data/Zebra-CoT-small subset. train supports one or more GPUs; submit GPU commands through the snapshot cluster launcher in docs/submission_snapshots.md. Runtime artifacts default to experiments/output/<run.dataset>/<run.category>/{sft,sft_probe}/<timestamp_name> unless an explicit output or SPATIAL_MMCOT_OUTPUT_ROOT is set. Probes always use smoke; failed or superseded runs are moved to the dataset-local archive only after they stop.

SFT publishes a step-zero recovery point, then checkpoints at configured update boundaries and every 30 wall-clock minutes. Re-executing the same snapshotted command in the same run directory restores the latest atomic checkpoint, discards only uncommitted log steps, and continues the exact logical schedule. This is the canonical recovery contract for SIQ-managed interactive jobs; final adapters are also atomic and idempotent.

The controlled one-epoch Zebra-CoT-small comparison uses config/sft/zebra_small_standard_1epoch.yaml (L_text + 10 L_bit) and config/sft/zebra_small_information_gain_1epoch.yaml (the same run plus both Information Gain directions). Both use exactly 1,800 train rows; the remaining 25 subset rows form one fixed validation set, with no experiment test split.

Information Gain uses one independent 2-D source-key mask per maximal adjacent block transition. Full M-RoPE positions, token IDs, image grids, target codes, and target prefixes are reused exactly. Text scores average target-token log probability; visual scores sum all 4x4x64=1024 Bernoulli log probabilities at each composite visual position and average only over AR positions. Blind scores are detached. Clean BSQ context and zero dropout are enforced while IG is on. The default samples one transition per trajectory and direction in a deterministic cycle; set event_selection: all to score every transition. IG always averages events within each trajectory first and then averages contributing trajectories, so trajectories with more transitions do not gain more objective weight. Validation pins the cycle index to keep checkpoint comparisons identical. When IG validation is enabled, each artifact also contains trajectory-balanced gain/absolute-gain and gate distributions, including saturation rates at the configured Temperature. Per-update logs record pre-clip gradient norm and the applied clipping scale. Full teacher-forced validation emits periodic rank-zero batch progress plus one completion line per rank, so long validation boundaries remain distinguishable from a stalled distributed job.

objective.bit.probability_mode defaults to fixed; uniform deterministically samples one rate in [0, input_flip_probability] per noisy trajectory and optimizer update before sampling its bit mask. IG's student_full_view_implementation defaults to fused, combining the differentiable noisy-SFT and clean-IG rows in one virtual-batch backbone call; sequential retains the prior two-call path.

objective.information_gain.formulation selects absolute (the backward- compatible default) or relative_per_unit. The relative path reuses the same full/blind forwards and direction settings, adds no model forward or material activation memory, and logs per-unit scores, dimensionless relative margins, and gates.

The full branch is computed once and shared by text CE, bit CE, optional Soft-BSQ features, and both full IG scores. Soft-BSQ reuses the hard target feature replay and adds only the differentiable frozen BSQ projections/mergers. Each event remains an independent blind view, but up to information_gain.blind_view_chunk_size views execute as one virtual batch. The default four-view chunk turns the two one-image views into one forward and the eight four-image views into two. All virtual rows stay in the four-layer visual decoder to preserve the verified FlashAttention2 batch path; only its final output projection is sparse. M-RoPE and frozen BSQ replay are shared across all three branches, and the 2,048-way visual output projection is evaluated only at supervised/scored carrier positions. Detached blind forwards disable autocast weight caching so the later differentiable output heads cannot reuse detached BF16 casts. Per-event score records are materialized only for validation/probes; the train path keeps sufficient statistics on device and avoids per-event host syncs. EMA-only clean students project only selected KD carriers. EMA shadow/native updates use cached references, preallocated buffers, and bounded foreach ops.

Atomic checkpoints preserve optimizer/scheduler, the exact data cursor, and per-rank RNG. Resume also verifies config plus cache-shard, base-model, and owned-source fingerprints. The packing budget is soft for an oversize singleton; data.max_sequence_tokens is the hard context limit. The default checkpoint.resume_mode: exact remains fail-closed. The exceptional ddp_execution_change mode requires a parent-linked run and permits only the two DDP reducer flags to change while restoring the same model, optimizer, scheduler, data cursor, and per-rank RNG. The parent-linked epoch_extension mode permits only a longer constant-LR epoch schedule and a new run name while keeping execution and objective settings exact. For single-GPU runs, optional trainer.execution_micro_batch_size and trainer.execution_batch_padded_token_budget re-chunk an already-fixed logical update without changing its rows, denominators, scheduler step, or BitFlip mask; the logical schedule hash is recorded in config.json. Distributed training defaults to native PyTorch DDP with backward/communication overlap. The loss compensates rank averaging only in DDP mode, so both modes implement the global numerator/count objective. ObjectiveOutput is a NamedTuple: its loss must be visible to DDP's pytree traversal and backward sink. An unregistered dataclass bypasses static-graph reduction in Torch 2.7. trainer.ddp_find_unused_parameters and trainer.ddp_static_graph expose the reducer graph policy. They default to true/false for historical configs; a workflow proven to cover every trainable parameter can use false/false to avoid rank-local reducer graph discovery. trainer.gradient_sync: ddp uses ordinary gradient buckets, with no custom communication hook or per-bucket CPU rendezvous. The v1 configs select find_unused_parameters=false, static_graph=false, and global batch 16 (4 GPUs x microbatch 4 x accumulation 1). Explicit post_backward SUM remains available for diagnostics and old-run comparisons. The data schedule and effective logical batch are unchanged; frozen jobs retain their source. Only gradient reduction uses the default NCCL group. Scalar statistics, finite-state flags, metadata, barriers, validation records, and checkpoint coordination use a separate CPU Gloo group so control collectives cannot interleave with asynchronous reducer buckets. Validation assigns whole shards to ranks, length-batches within each shard, and gathers event records only to rank zero. All ranks compare full trainable-gradient hashes after the first backward of every launch and full parameter hashes before checkpoint writes; replica drift fails before committing an update or checkpoint. Parameter schemas and required text/visual output gradients are also checked. SFT requires Torch >=2.7.1,<2.8 and uses default process-group options with eager device_id. Torch 2.7.1 fixes the implicit nonblocking-mode regression in 2.7.0 (PyTorch #153960). Leave NCCL API-mode environment overrides unset; CUDA stream asynchrony and watchdog/error monitoring remain enabled. trainer.process_group_timeout_minutes and trainer.nccl_heartbeat_timeout_seconds separately configure the NCCL collective timeout and watchdog heartbeat; both default to 20 minutes while monitoring remains enabled.

On-policy distillation

src/opd/ is separate from SFT and imports reusable logic only from src/shared/. It collects a fresh deterministic or temperature-sampled hard text/BSQ trajectory from $B_0$ for every logical update and verifies the full serialization once. The historical selector scores the first complete generated-image-to-text window; all_complete_transitions scores every T→V/V→T event with trajectory-first symmetric JSD. The student owns <|vision_start|> and the controller owns only <|vision_end|>; every other scaffold token and <|im_end|> remains an autonomous target. Three routing modes and none/uniform/failure GT replay form eight valid combinations; failure replay is available for both correct-only and privileged-correction. The historical seven Visual-Jigsaw configs explicitly share the canonical one-epoch half-BitFlip vanilla-SFT adapter through model.initial_adapter; OPD checkpoint resume is a separate field and restores optimizer, scheduler, teacher, cursor, and RNG. In-place retries prefer their own latest complete checkpoint over an explicit parent checkpoint in the frozen argv. OPD_RESUME reports that decision. --stop-after-updates N commits and pauses after N new updates for restart tests without changing the configured training horizon. The optional rollout.max_generated_tokens caps the entire continuation, including visual positions and scaffold but excluding the prompt. Unfinished capped rollouts count as wrong and enter fixed-budget GT failure replay. Use rollout_limit_change for an explicitly requested cap change from a parent checkpoint; it retains training state and refreshes the validation baseline.

PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m opd inspect \
  --config OURS/config/opd/zebra_visual_jigsaw_opd_all_none.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m opd train \
  --config OURS/config/opd/zebra_visual_jigsaw_opd_all_none.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR torchrun --standalone \
  --nproc-per-node=4 -m opd train \
  --config OURS/config/opd/v2/thinkmorph_privileged_failure_c11c10_base.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR torchrun --standalone \
  --nproc-per-node=4 -m opd train \
  --config OURS/config/opd/v3/zebra_visual_jigsaw_ig_abs3e_bidirectional_base.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m opd audit-rgb \
  --config OURS/config/opd/zebra_visual_jigsaw_opd_all_none.yaml \
  --samples 1 --decoder-steps 28

OPD v0 supports one or multiple GPUs in one node and executes one sampled trajectory per rank-local scoring chunk. Ranks cost-balance one logical rollout/replay budget by prepared sequence length, then SUM globally normalized gradients once per update. Necessity-only runs fuse C11/C10 teacher masks into one no-grad virtual batch. Rank 0 alone verifies training rollouts and writes training artifacts. Validation now drains one shared queue across all GPUs and evaluation.workers_per_gpu independent workers per GPU (A100/A40: 1; GH200/H200: 2). Atomic per-sample results resume within the same validation step. Training also supports rollout.workers_per_gpu with the same persistent pool; all replicas receive one current student snapshot before each logical batch. Validation saves full hard trajectories and all numeric metrics to TensorBoard. Pre-clipping OPD/replay gradient norms, ratio and cosine are measured every logging.gradient_diagnostics_every_updates updates without extra backward. The v3 config enables teacher top-K union with a per-position retained-mass guard and text contribution clipping; older configs retain full-vocabulary JSD. The optional deepseek_answer verifier first applies deterministic answer parsing and calls deepseek-v4-flash only for unresolved semantic matches. Visual OPD averages Bernoulli JSD over the 1024 BSQ factors before reducing positions; modality weights therefore do not compensate for factor count. The ThinkMorph base uses a declarative answer-only teacher privilege and warmup followed by constant LR. Its text categorical and visual independent-Bernoulli sampling both use temperature 0.7; stateless sample/update seeds make retries exact and give paired validation common random numbers. Canonical rollout limits are null: native context length is the hard horizon, with shared 3-repeat/period-at-most-4 cycle detection; explicit positive limits remain available only for smoke/debug runs. Artifacts retain complete sampled token IDs, hard BSQ codes, ownership, grids, verifier descriptor/result, routing, replay selection, diagnostics, and atomic student/teacher/optimizer/scheduler/RNG checkpoints.

On-policy reinforcement learning

src/rl/ implements config-driven one- or multi-GPU multimodal GRPO v0 independently of OPD. Each optimizer update collects fresh i.i.d. stochastic full-trajectory prompt groups, retains invalid, wrong, and truncated samples, computes per-prompt population-normalized advantages, and performs exactly one actor update. The strict verifier currently accepts only Visual-Jigsaw A–D references. Image reward is always literal zero, does not decode RGB, and has no backend.

trainer.world_size: 1 is the canonical single-GPU path. A larger world size may not exceed the rollout group size; the unchanged CLI entry point auto-spawns that many local CUDA workers. RLv3 workers dynamically claim prompt/group members with rank-independent seeds; static round-robin is available for comparison. They all-gather the fixed group, compute the same global rewards/advantages, sum gradients already divided by the global trajectory count, and synchronize clipping, optimizer, scheduler, checkpoint, and validation state. Uneven cases such as $G=4$, world size 3 use a deterministic 2/1/1 split in static mode without changing the objective. Rank 0 writes update/checkpoint artifacts; validation workers atomically commit their own query bundles and claim unfinished queries from a shared queue. Validation normally reuses the training rollout multiplicity; an explicit evaluation.rollouts_per_sample instead scores that many independent trajectories per sample without constructing GRPO advantages. This permits a single-rollout validation pass even when training uses a larger group. RL uses the same context-bound default horizon and shared cycle detector as OPD; explicit positive text/phase/image/visual limits are smoke/debug controls.

The actor uses a uniform autonomous-AR-step reduction with factor-wise visual PPO clipping. A frozen run-start actor supplies full-distribution symmetric JSD regularization. The fresh one-update path uses a sampling-old-value-aligned, full-prefix-gradient score-function surrogate; sampled old, aligned surrogate, raw replay, reference, and distributional mismatch diagnostics remain separate in artifacts. Replay means retain their per-trajectory safety bound, while percentile guards reduce the complete logical-batch action population rather than taking the maximum of trajectory-local quantiles. This path is not reusable as multi-update PPO. Generated grid sides lie in 1–30 and are parsed only after the student samples <|vision_start|>, so termination adds no hidden policy decision. Optional IG uses detached pre-update behavior-policy $11/10/01/00$ self-scores. Source-copy is measured online; blank/gray/noise, OCR/text-overlay, reward-shadow, and BSQ↔RGB diagnostics are only marked for offline/shadow evaluation of saved rollouts.

PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m rl inspect \
  --config OURS/config/rl/zebra_visual_jigsaw_grpo_fmt_acc.yaml
PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR python -m rl train \
  --config OURS/config/rl/zebra_visual_jigsaw_grpo_fmt_acc.yaml

Atomic checkpoints preserve actor, optimizer, scheduler, frozen-reference descriptor/state, query/update cursors, and RNG. Checkpoints link committed rollout artifact directories and hash their complete file inventory; only committed attempt metrics enter global JSONL/TensorBoard, while uncommitted retries remain forensic only. Validation streams aggregates and reuses results only under an exact actor/config/sampling/sample-set identity match. The validation identity separately records training group size and evaluation rollouts per sample. Resumed runs sample a fresh group; trajectories are never carried across updates or stored in checkpoints.

Verify

PYTHONPATH="$PWD/OURS/src" mamba run -n UniAR pytest -q OURS/tests

Historical workflows and their artifacts live under experiments/. When running them without installing this package, place the backbone first:

export PYTHONPATH="$PWD/OURS/src:$PWD/experiments/src"