docs/thinkmorph.md

ThinkMorph workflows

Canonical BAGEL-family framework

OURS-Bagel/ owns new training work for the shared BAGEL architecture, including both models/Bagel-7B-MoT and models/ThinkMorph-7B. SFT v1 is implemented with strict YAML, a typed configurable interleaved protocol, protocol-fingerprinted Parquet conversion, native full-parameter FSDP, global-count-normalized CE/flow-MSE, EMA, atomic resume, and per-update TensorBoard/JSONL. LoRA opd implements the text/flow contract below; rl is reserved.

The first canonical configs are OURS-Bagel/config/sft/thinkmorph_paper_bagel7b.yaml and the one-update thinkmorph7b_sft_smoke.yaml. New jobs use cluster_jobs.submit --source-profile=bagel; outputs live below experiments/output/<base-model>/<dataset>/<version>/...; see docs/artifacts.md.

BAGEL SFT supports two explicit controls. updates preserves ThinkMorph's 32K token-packed, optimizer-step recipe. epochs uses fixed source-sample packs and derives the update target from dataset size and logical batch. The first rank-128 LoRA candidate uses the full original trainable component scope, constant LR 2e-5, four samples per rank on four GPUs, logical batch 16, and ten sample-counted epochs. Its full parity audit is docs/research/20260914-225604_bagel_thinkmorph_lora_sft_parity_audit.md.

The main candidate uses data/ThinkMorph/Bagel-SFT-IDVal5, which projects the audited source-coordinate split onto the native BAGEL Parquet without changing images or rendered text. It contains 23,741 train rows and 1,249 validation rows (5%). Epoch mode shuffles every train row deterministically, pads only the three-slot final logical batch, and resumes by update number. Validation uses the EMA model, disables all condition dropout, excludes the 15 distributed tail-padding rows from losses/counts, and runs at step 0 and every epoch. The BAGEL Qwen stack routes packed-training and query-inference calls by their explicit keyword contracts rather than Module.training; EMA can therefore remain in deterministic eval mode while using the loss-producing packed path. Rank-128 LoRA uses replicated DDP, not FSDP: every rank constructs a zero-data meta skeleton, rank zero reads the base once, and DDP broadcasts it; every GPU stores one BF16 frozen backbone and FP32 LoRA parameters. EMA is an FP32 adapter shadow temporarily swapped in for validation, so no second backbone is constructed. Full-parameter paper reproduction remains FSDP. Checkpoint assign restores persistent tensors; the omitted non-persistent rotary buffer is rebuilt before moving rank zero to GPU, while other ranks use to_empty and receive the synchronized state. Condition dropout retains its original behavior. DDP uses find_unused_parameters=true with a dynamic graph; the attempted static-graph optimization was removed after it triggered a reducer assertion during gradient accumulation. LoRA DDP uses activation checkpointing: disabling it reached roughly 137.5/139.8 GiB even at smaller execution packs and exhausted memory again when Adam state was initialized. The stable execution is four samples/rank with accumulation one. If a rank drops every ViT carrier, BAGEL skips the ViT branch for that batch; dynamic DDP then treats its ViT LoRA parameters as legitimately unused. One-hour Delta SIQ attempts write atomic recovery checkpoints every 100 updates and at every validation boundary; wall-clock checkpointing is disabled. The current step-17599 continuation runs on DeltaAI GH200×2 with microbatch 2 × accumulation 4, preserving logical batch 16 and the fixed logical-batch sample cursor. It restores adapter, EMA, AdamW, scheduler and both rank RNG states. Subsequent TIMEOUT retries prefer the continuation's own newest complete checkpoint. Attempts use the two-hour ghx4-interactive queue. Its config is OURS-Bagel/config/sft/thinkmorph_lora_r128_constant_val5_16epoch_2gpu_resume17600.yaml; it resumes the four-GPU continuation's latest complete checkpoint 0017599 and continues toward 23,744 total updates, retaining the parent TensorBoard history. After the A6000 trial OOMed before publishing a checkpoint, the original frozen GH200×2 run was re-enqueued in the same output directory. Its newest committed checkpoint is 0019599; it resumes at update 19,600 with GPU-resident EMA. The 16-epoch continuation config OURS-Bagel/config/sft/thinkmorph_lora_r128_constant_val5_16epoch_4gpu_continue.yaml extends the completed 10-epoch checkpoint 0014839 to 23,744 total updates. It uses four GH200s, two samples per rank and two accumulation microsteps. Epoch extension requires a completed constant-LR parent at its final update; seed and all mathematical settings stay identical. Explicit hardware migrations may adjust processes, samples per pack and accumulation while preserving logical batch 16; added ranks use their seeded RNG when no parent-rank RNG exists. It restores student/EMA/optimizer/scheduler/RNG state, inherits TensorBoard history, and writes new checkpoints in its own parent-linked run directory.

Lab continuation uses /mnt/ultracube/tac010/MMCoT, shared by simiao and luogeng, and the private shared environment /mnt/ultracube/tac010/apps/miniforge3/envs/ThinkMorph. The A6000×8 config thinkmorph_lora_r128_constant_val5_16epoch_luogeng_a6000x8.yaml resumes checkpoint 0019599, with train microbatch one and accumulation two (logical batch 16). Validation uses one sample per rank while retaining all 1,249 rows and padding masks. distributed.ema_cpu_offload=true keeps the FP32 EMA shadow and validation student backup on CPU; updates retain the same EMA equation and checkpoint keys. NCCL loading stages one adapter tensor at a time on CUDA. This frees about 2.7 GiB per rank; GPU-resident EMA at microbatch one reached 47 GiB and OOMed on the first A6000 batch. CPU EMA passed the first updates but later also OOMed, so this configuration is not validated for the full run. Hardware continuation may change validation pack size and EMA placement. Transfer only the BAGEL source closure, base weights, IDVal5 cache, selected checkpoint and parent metadata/TensorBoard through the EmergentTwilight/MMCoT bucket. thinkmorph_sft.relocate_bagel_cache rebases copied JSON path metadata and retains its original files. Prepare a parent-linked source snapshot on the lab filesystem and execute its frozen script through luogeng's gpuq with eight exclusive GPUs; Slurm directives are inert when that script runs in Bash.

Independent OPD replicas stage the originating run's immutable source closure under experiments/cache/ through a task-scoped HF bucket prefix, leaving the lab's live source tree unchanged. Prepare a BAGEL snapshot from that staged source project and record env SPATIAL_MMCOT_PROJECT_ROOT=/mnt/ultracube/tac010/MMCoT in the launch argv, so models/data still resolve from the canonical lab project while SPATIAL_MMCOT_SOURCE_ROOT remains the new frozen source. Use the explicit /mnt/ultracube/tac010/apps/miniforge3/envs/ThinkMorph prefix: a separate empty /mnt/ultracube/tac010/envs/ThinkMorph can shadow a name-only mamba invocation. Luogeng's approved http://webproxy.ucsd.edu:3128 HTTP(S) proxy is required for Hub/judge access; pin it per task, without changing scheduler/global settings. Transfer judge credentials only over SSH into the ignored, mode-0600 key file; keep them out of buckets, snapshots, argv and logs.

A same-start replica copies the original step-zero checkpoint (including student/EMA/reference and all four rank RNG states), not the active run's newest checkpoint. lab_replica.json records source provenance, GPU allowlist and the future gpuq argv. prepared_not_queued means no GPU task has been submitted; only an explicit later launch should execute that argv. Four-GPU replicas keep logical batch 16 and use --gpu-count 4 --gpus 0,1,2,3; the queue owns device visibility. The cluster run is independent and remains untouched.

For a two-GPU launch of an unstarted replica, use config/opd/thinkmorph_igopd_nec_all_2gpu.yaml; only trainer.processes changes from four to two. The dynamic queues still collect 16 fresh trajectories per optimizer update and backpropagate one trajectory at a time, so no separate accumulation or batch-size change is needed. Prepare a new parent-linked snapshot and invoke thinkmorph_sft.prepare_opd_initial_world_change before queueing. It rejects every config difference except process count, rejects trained/Adam-updated checkpoints, preserves student/EMA/reference/optimizer/ scheduler tensors, and copies parent RNG ranks 0/1 into the two-rank checkpoint. Only the checkpoint world/config identity is rewritten, with provenance in initial_world_change.json. This initializer does not relax ordinary runtime resume checks or support world-size changes after training has started. Use gpuq --gpu-count 2 --gpus 0,1; retain the unstarted four-GPU preparation.

If this two-GPU launch is stopped during initial validation before any optimizer update, the untouched four-GPU preparation can be activated again. Before reusing validation bundles, hold writer locks, require both latest checkpoints to be step zero, compare every checkpoint field/tensor except world/config identity, verify matching core-source hashes and config fields except process count, and validate each sample/seed identity. Copy bundles and judge diagnostics only into an absent target validation directory; preserve the source and the original four-rank RNG files. Record this cache import in validation_reuse_from_2gpu.json. This is reuse under an identical initial policy, not permission to rewind a trained run or relax runtime resume guards.

projects/ThinkMorph-edited supplies pinned architecture, packer and FSDP execution code. It accepts an external dataset registry and both EMA-only Bagel-7B and model-only ThinkMorph-7B initialization. Experiment configuration, protocols, preparation, launch and artifact publication belong only to OURS-Bagel.

Canonical BAGEL on-policy distillation

bagel_mmcot.opd implements docs/modeling.md §4. The strict schema and entry point are config/opd/thinkmorph_joint.yaml and python -m bagel_mmcot opd train --config ... --output-dir .... Default initialization is models/ThinkMorph-7B/model.safetensors; original BAGEL uses model.path=models/Bagel-7B-MoT, model.weights=ema. An optional canonical SFT adapter may initialize the student; its rank/alpha and target scope must match. LoRA is FP32 over a BF16 backbone and the VAE stays frozen.

• Text targets use exact sampled IDs, including model-owned EOS. Full-vocabulary JSD, forward KL and reverse KL are selectable; positive log-ratio clipping shapes the detached EMA target. Vocabulary projection is checkpointed by target-token chunk rather than dropping support mass.
• Flow targets use detached (x_t,t) recorded before selected steps of the student's actual reverse Euler solve. Selection is uniform without replacement and seeded independently of the solver. Targets use signed, unclipped teacher velocity extrapolation and coordinate-mean MSE. v1 validates raw data-to-noise velocity, Euler and CFG scales one explicitly; guided inference fields and alternative solvers are not implemented.
• All four views preserve query carriers, block layout and RoPE positions. Scoring reconstructs the complete packed prefix; hidden text payload keys and both VAE/ViT generated-image payloads are zeroed and masked at every layer. Structural marker pieces remain visible. Frozen VAE posterior samples use a content-derived seed/cache, shared by rollout and all scoring views. Learned vision/projection features are recomputed with the active parameter view.
• Necessity uses student 11; sufficiency uses student 01 and excludes events without earlier generated blocks. Components reduce event means within a trajectory, then means over globally accepted trajectories with that component.
• Reference is the run-start student adapter, persisted independently of EMA. It evaluates accepted event unions once, always full-visible and unprivileged, with the same sampled prefix/flow state. Direction means are equally weighted over the fixed enabled directions. Reference weight zero disables this term.
• Routing modes are all, correct_only, privileged_wrong. Only paired teacher views receive answer privilege; student and reference never do. Verifiers use normalized exact answers or rule-prefetched DeepSeek semantic equivalence. Presentation-only outer answer/math/boxed wrappers and comma spacing normalize locally; leading decimal points remain significant. This does not infer an answer from full reasoning, inspect images, or simulate navigation paths. API errors remain errors, never silently converted to incorrect labels.
• Replay modes are off, uniform, failure_first. Its fixed global budget uses current failed indices first when requested, then samples outside all current raw indices. Replay rereads complete native GT, with unweighted text CE and flow MSE globally normalized by their own target counts. Condition dropout is disabled. Its logit-normal forward noising is separate from OPD's saved reverse states.

OPD uses two dynamic queues per frozen logical batch:

1. GPU workers atomically claim prompts and commit each trajectory/verdict.
2. After every logical slot completes, determine routing, component/trajectory denominators and fixed-budget GT replay in logical-slot order. Native replay packs are counted before backward; their deterministic reconstruction is checked against those counts.
3. One shared backward queue contains whole trajectory-scoring tasks and whole GT-replay tasks. Any optimizer rank can read another worker's committed trajectory and claim the next task. Event/query losses are backpropagated serially within a trajectory to bound activation memory.
4. Only after every backward task completes, SUM globally normalized gradients across bounded buckets, clip once, apply AdamW once, then update EMA once.

trainer.processes is the number of GPUs/optimizer ranks. Logical batch remains 16 in the canonical config; it need not divide the rank/worker count, and a rank may receive zero tasks. rollout.workers_per_gpu and evaluation.workers_per_gpu include the main rank and default to one. Values above one spawn persistent inference-only replicas on that rank's assigned GPU; their student adapters refresh through shared CPU buffers before each phase. They do not join NCCL or own optimizer/EMA/reference state. Backward uses one replica per GPU to bound memory, with dynamic ownership across GPUs. Extra inference replicas retain their backbone allocations during backward: worker counts must fit the combined peak, not just standalone inference memory.

Python, NumPy and Torch randomness is scoped to sample/update/phase rather than worker identity, including the native GT packer's logit-normal time draws. Queue completion order cannot alter failure-first selection or normalization. Floating-point accumulation order can still differ; numerical agreement, not cross-hardware bitwise identity, is the reproducibility contract.

Queue claims are attempt-local FileStore state. Workers wait on file-store completion/failure markers before entering collectives, so a healthy long query does not consume a collective timeout. scheduling/ exposes the latest attempt's task owners/timings in logical order; _runtime/<attempt>/audit/ retains each attempt's original scheduling evidence. Reference remains frozen. Losses, coverage, clipping and fixed validation go to TensorBoard and per-update JSON. Each sample first atomically saves text IDs, images and solver states with verdict: null, then atomically finalizes the same bundle after valid judging. Retries reuse that pending trajectory and rerun only the judge; committed verdicts remain immutable. A pending/unverified trajectory cannot enter backward. Complete checkpoints include optimizer/scheduler, student, EMA, reference identity/weights and rank RNG; an initial checkpoint precedes rollout. checkpoint_every=1 is required: every reusable query batch must have a committed behavior policy. Recomputing several uncheckpointed updates under a different task order could change floating-point weights and invalidate later saved trajectories, even when their update numbers/seeds match. Identical argv resumes the newest complete in-run checkpoint, rebuilds fresh queues and reuses committed queries with matching behavior update/sample/seed. An interrupted backward discards partial gradients and reruns the whole uncommitted update; it never skips tasks because a previous attempt claimed them. Scheduled validation interrupted after its checkpoint is completed before the next training batch (including validation at the final update).

New experiments use submission-time BAGEL snapshots and SIQ for interactive partitions. config/opd/thinkmorph_joint_a40_smoke.yaml reduces rank, images, solver/query budget and logical batch for a runtime check. The wrapper thinkmorph_sft.opd_resume_smoke invokes identical argv twice and verifies reference immutability, optimizer progression and nonzero native components. GPU execution evidence is recorded in its smoke_report.json; successful CPU tests alone do not establish GPU or learning quality. thinkmorph_dynamic_a40_smoke.yaml runs two optimizer ranks with a smaller logical batch for queue verification. The same wrapper's --interrupt-collection option deliberately terminates only its own process group after a committed query, then checks identical-argv query reuse, distinct backward tasks and checkpoint progression. The deterministic CPU suite also checks skewed queues, peer failure, persistent worker refresh, and logical-16 distributed-versus-serial gradient agreement. The A40 smoke 20260929-162430_joint_text_flow_nec_suf_reference_replay_a40 completed two updates/four fresh rollouts across two identical launches. It restored EMA counter 1 and immutable reference counter 0, then verified the reference tensors remained exactly equal at update 2. Text nec/suf and flow nec were nonzero; pre-clipping norms were 0.02385 and 0.47283, with no clipping and 29.27 GiB peak allocation. Each sampled trajectory had one generated image, so T→V sufficiency was correctly absent; its multi-image masks/gradient/query semantics have CPU regression coverage, not real multi-image GPU evidence yet. All four answers matched the exact verifier, so this smoke exercised replay's batch-external fallback; incorrect/invalid routing and failure-first selection are covered separately by deterministic tests. These are runtime checks, not evidence of accuracy or image-quality gains.

Necessity-only routing comparison and memory smoke

config/opd/thinkmorph_igopd_nec_all.yaml and thinkmorph_igopd_nec_privileged_wrong.yaml differ only in name and routing. Both disable sufficiency without renormalizing the remaining weights: T→V and V→T necessity each have weight 0.25; guidance alpha is respectively 0.25/0.5. They use LoRA-r16, LR 2e-6, logical batch 16, uniform replay budget 4 with weight 0.01, frozen-reference weight 0.1 and EMA decay 0.995. all admits every structurally valid complete trajectory, including wrong answers; privileged_wrong adds the verified GT answer only to wrong trajectories' paired teacher views. Invalid/truncated trajectories enter neither OPD branch. Student, reference, and validation rollout never receive answer privilege.

The text cap is 4,096 new tokens per text block, not per trajectory; max_images=5 counts generated images. The separate 16,384-position context guard remains in place, so these are independent stopping bounds, not a promise to generate five images plus six full-length text blocks. Images retain native 512–1,024 geometry, and raw Euler sampling retains 49 intervals/four queries. Every 16 updates, the current student autonomously rolls out the same fixed 128 held-out queries and the verifier checks final answers. This is rollout accuracy/validity, not teacher-forced token accuracy or validation loss. Current validation inherits the training rollout temperature/sampling/limits; there is no separate greedy-validation switch. Step-zero and final validation also run; identical-argv retries reuse committed validation samples.

For the A100×4 30-minute memory check, thinkmorph_igopd_nec_a100_smoke.yaml changes only the run name, horizon to one update, validation count to four and trainer.memory_profile=true from the privileged configuration. It retains logical16, r16, full image resolution, solver budget and long rollout caps. The entrypoint python -m thinkmorph_sft.opd_memory_smoke --config ... --output-dir ... reports all ranks' allocation/reservation peaks, phase errors, actual generated length/image counts and which necessity components executed. Profiling records live state sizes and CUDA peaks under memory/rank-<rank>-<attempt>/; it is disabled in production and preserves lifetime peak metrics when enabled. Natural short-rollout success does not establish worst-case 4,096-token/five-image safety. A timeout or an unexercised branch is not reported as a passed memory test.

Robust judging and startup recovery

The DeepSeek fallback puts the equivalence rule in a system message and serializes question/reference/candidate as JSON data. It requests json_object output with one boolean correct field and explicitly disables thinking. The output budget starts at 256 tokens and doubles to at most 1,024 across the existing bounded retry count. Empty/malformed/truncated responses are retried with backoff; HTTP 408/429/5xx and connection failures are retryable, while permanent client errors (such as authentication) fail promptly. Only finish_reason=stop can yield a verdict. Exact, full-string legacy labels such as A., (A) or Answer: B are accepted, but prose containing letters, duplicate JSON keys, string/number booleans and ambiguous answers are rejected.

Every API attempt records the judge protocol, sample, request data, response content (bounded/redacted), finish reason, usage and error class in the sample's adjacent <slot>.judge/ directory. Authorization headers are never recorded; the active API key and key-like strings are redacted. Exhausted retries still abort the uncommitted update, preserving the pending rollout for the next launch. The archived failure's original API response was not saved, so its exact text cannot be reconstructed; formatting and pending-rollout regressions exercise the known failure paths without claiming to reproduce that lost response.

For a verifier failure before the first optimizer update, prepare a distinct source snapshot with --parent-run, then run python -m thinkmorph_sft.prepare_opd_verifier_retry --parent-run ... --run-dir ... --config ... before SIQ enqueue. This requires an unchanged full config/world/reference identity, a step-zero checkpoint whose student/EMA/reference are equal, no Adam update state, complete rank RNG files and an inactive parent writer. It copies that state byte-for-byte and validates each cached query's split/slot/sample/ seed/behavior identity. Generated trajectories are reused as pending; old judge decisions and validation aggregates are deliberately not carried over to the repaired judge protocol. verifier_retry.json records the lineage and selected samples. The parent remains untouched. Later SIQ TIMEOUT retries use the child's ordinary newest complete checkpoint and committed query cache. This narrow startup repair does not authorize changing a trained run's judge mid-experiment.

Each GPU stores one frozen BF16 backbone, student FP32 LoRA and its gradients/ Adam states. EMA and reference have independent FP32 LoRA shadows on CPU; a no-grad teacher/reference call temporarily loads that shadow and restores the student afterward. No second full teacher backbone or teacher autograd graph is retained. For the native ThinkMorph-7B checkpoint and scope=mot,r=16, header-derived estimates are: backbone 27.21 GiB, VAE 0.16 GiB, student LoRA 0.30 GiB, gradients 0.30 GiB, Adam moments 0.60 GiB (about 28.6 GiB fixed GPU storage), plus one temporary 0.30 GiB adapter view. Each CPU shadow is 0.30 GiB. Activations, KV caches, vocabulary chunks, attention masks, NCCL and allocator workspace are additional. These are estimates, not measured peaks; A100×4 replicates this footprint on each card instead of pooling four cards' memory.

Legacy native ThinkMorph-7B OPD prototype

ThinkMorph-OPD/ contains the earlier native models/ThinkMorph-7B backbone, protocol tracer, and V→T-only LoRA OPD trainer. It imports BAGEL/MoT architecture code from projects/ThinkMorph-edited but never substitutes models/Bagel-7B-MoT weights. The first real trace, exact token IDs, heterogeneous VAE/ViT/diffusion carriers, controller state machine, and LoRA boundary are recorded in docs/research/20260910-222300_thinkmorph7b_native_token_and_state_machine.md. It is retained for evidence and migration tests; new OPD changes wait for its audited move into OURS-Bagel.

The first native objective selects every complete generated-image→next-text block and implements Necessity only. EMA C11 includes both the generated clean-VAE and ViT carriers; EMA C10 creates shape-matched carriers without encoding RGB and marks every VAE/ViT payload KV invisible to target attention. It preserves <|vision_start|>, <|vision_end|>, logical carrier length and the two image RoPE steps. Both views receive the same answer privilege only for an incorrect rollout. They score the unchanged sampled text block, including its ChatML end target. A positive C11−C10 log-ratio shapes the C11 teacher target, which is distilled to the full-visible student with symmetric JSD.

Scoring reconstructs historical carriers under stop-gradient, then retains the student graph through the current generated image and target block. The base adapter covers understanding and _moe_gen attention/MLP linears; SigLIP, VAE, connectors and output heads stay frozen. Its support is the deduplicated union of C11/C10 top-128 tokens, with retained-mass checks and the same per-coordinate 0.05 clipping semantics as the public Visual-OPSD implementation.

The DeepSeek verifier is rule-prefetched and fail-closed. Current-update wrong or invalid sample indices fill the fixed replay budget first; any deficit is uniformly filled outside the raw rollout batch. Replay re-reads the independent GT sample and reconstructs the original ThinkMorph training plan: observation VAE+ViT, instruction, GT text-1, noised GT-image rectified-flow target, the same image as clean VAE+ViT context, and GT text-2. It applies length-reweighted text CE plus visual-position-mean flow MSE, and never pairs a wrong sampled image with GT continuation. T→V remains absent from OPD but is preserved by replay.

configs/v2t_opd_base.yaml uses logical batch 16 on four GH200 ranks, 128 updates (2,048 fresh rollouts), rank/alpha 16, rollout temperature 0.75, LR 5e-6, EMA 0.995, JSD β/temperature 0.5/1.0, contrast alpha 0.5, and replay 4/16 with outer coefficient 0.01 and text/flow weights 1/1. These values and their evidence status are audited in docs/research/20260910-225127_thinkmorph_native_v2t_opd_hyperparameters.md. After this 2,048-rollout pilot passes its stop criteria, configs/v2t_opd_main.yaml is the recommended main experiment: 1,000 updates and 16,000 unique task-proportional rollouts, retaining all objective/replay settings while using the official-quality 50-step flow sampler, 50-update warmup, validation every 50 updates and recovery checkpoints every 16 updates. The trainer writes every rollout atomically, synchronizes globally normalized LoRA gradients with one flattened collective per dtype instead of hundreds of per-tensor collectives, checkpoints optimizer/scheduler/EMA state, and resumes in the same prepared run directory.

The revised GH200×4 smoke completed 4 updates/16 rollouts in 5m06s, including attention-masked C10, both-expert LoRA/EMA, full text+flow replay, flattened gradient SUM, start/end validation, per-update checkpoints, TensorBoard and student/EMA adapters. Peak allocated memory was 79.19 GiB and all gradient norms were finite without clipping. Its replay coefficient 0.1 was numerically too large relative to raw OPD, so canonical pilot/main configs use 0.01. No pilot or main experiment was submitted.

Paper-aligned BAGEL baseline

The execution fork is projects/ThinkMorph-edited, pinned from upstream commit c1a48ad; OURS-Bagel owns the baseline and future BAGEL objectives. This is separate from the UniAR conversion below. The starting model is the pinned models/Bagel-7B-MoT EMA checkpoint, and its isolated DeltaAI environment is ThinkMorph (Torch 2.7.1+cu128, Transformers 4.49.0, FlashAttention 2.8.3.post1).

data.convert_thinkmorph converts all four released datasets to the official three-column BAGEL contract under data/ThinkMorph/Bagel-SFT. It preserves the official system prompt and tag strings byte-for-byte, re-encodes both images as PNG, writes 100 rows/file and 10 rows/row-group, and publishes a hash manifest, parquet_info.json, and READY atomically.

OURS-Bagel/config/sft/thinkmorph_paper_bagel7b.yaml is the canonical four-GH200 baseline: 8,000 optimizer updates, LR 1e-5, cosine with 2,000-step warmup, CE/MSE weights 1, condition dropout 0.1/0.3/0.3, 32K carriers per rank microstep, four microsteps per update, BF16 four-way FULL_SHARD, activation checkpointing, full-parameter AdamW, EMA 0.9999, and clip norm 1. The accumulation window uses one global CE denominator and one global MSE denominator. The smoke entrypoint keeps the 8,000-step scheduler but stops after 20 committed updates.

The primary-source audit and the known defects in the public launch script are recorded in docs/research/20260905-140823_thinkmorph_official_sft_reproduction_contract.md.

DeltaAI environment

MAMBA=/work/nvme/bgjw/tcheng2/apps/miniforge3-deltaai/bin/mamba
$MAMBA create -n ThinkMorph python=3.12 pip setuptools wheel ninja -y
$MAMBA run -n ThinkMorph python -m pip install 'torch==2.7.1+cu128' \
  --index-url https://download.pytorch.org/whl/cu128
$MAMBA run -n ThinkMorph python -m pip install --no-deps \
  'torchvision==0.22.1' --index-url https://download.pytorch.org/whl/cu128
$MAMBA run -n ThinkMorph python -m pip install \
  -r projects/ThinkMorph-edited/requirements-gh200.txt

Install FlashAttention 2.8.3.post1 from an SM90/Torch-2.7-compatible wheel or build it on a GH200 using the CUDA 12.9 procedure in docs/uniar.md. The current isolated environment passes pip check and imports Torch 2.7.1+cu128, Transformers 4.49.0, and FlashAttention 2.8.3.post1.

The matching x86 environment on Delta is /work/nvme/bgjw/tcheng2/apps/miniforge-delta/envs/ThinkMorph; it has the same Torch, torchvision, Transformers, and FlashAttention versions. Do not reuse the DeltaAI AArch64 prefix on Delta.

VLMEvalKit UniAR evaluation

projects/VLMEvalKit_Thinkmorph_llm is the canonical evaluation fork and has a separate ThinkMorph-Eval environment under both Miniforge installations. It uses Torch 2.7.1+cu128, Transformers 4.57.0, PEFT 0.20.0, FlashAttention 2.8.3.post1, and the fork's evaluation requirements. pip check and the full VLMEvalKit configuration import must pass on both architectures.

The registered model keys are uniar-sft, the rank-64 two-/four-epoch development adapters, the two-epoch BitCE weight-1/weight-3 ablations, and the two-epoch absolute-IG adapter uniar-thinkmorph-v1-ig-abs-r64-val5-2epoch. The adapter injects only the canonical training protocol as the system prompt. It preserves the VLMEvalKit image/text user message exactly, with no appended instruction, uses the canonical full-prefix interleaved controller, and returns the readable assistant sequence. Set SPLIT_THINK=1 so MCQ scoring receives only the post-thought final segment. VSP is the exception: inference retains the full assistant response even when SPLIT_THINK=1, because its upstream scorer extracts the boxed path from the reasoning trace. Prompt-controlled VSP studies use the registered VSP_maze_task_main_original_60_balanced dataset (ten deterministic examples per maze size 3--8). Model keys ending in vsp-thinkmorph-short or vsp-short replace only the framework-owned VSP user content with the exact training prompt and its final test image; they retain the UniAR system protocol, chat template, generation path, and upstream strict scorer. VSP scoring defaults to the fork-compatible thinkmorph_strict mode. Set VSP_VALIDATION_MODE=official to reproduce the official FrozenLake transition semantics: off-grid actions leave the state unchanged and reaching the goal terminates successfully. Mode-qualified result filenames prevent cache reuse across scoring definitions. Native models/ThinkMorph-7B uses thinkmorph-7b-vsp-short-a40 for the same short-prompt VSP protocol; its A40-safe model-side batch is four. BLINK-Jigsaw is cached at data/VLMEvalKit/BLINK_Jigsaw.tsv; its MD5 is 174fc5d2ed0d67f9e0b33cb8a8be6727 and it contains 150 rows.

Canonical LLM-assisted evaluation uses the fork's original verifier with DeepSeek V4 Flash as its fallback choice extractor:

LMUData=data/VLMEvalKit SPLIT_THINK=1 torchrun --standalone \
  --nproc-per-node=4 projects/VLMEvalKit_Thinkmorph_llm/run.py \
  --data BLINK_Jigsaw \
  --model uniar-thinkmorph-v1-vanilla-r64-val5-2epoch \
  --judge deepseek-v4-flash --reuse --work-dir <prepared-run>

The completed canonical DeepSeek-V4-Flash-judged results supersede all scores from the retired local exact matcher:

Epoch	Bit CE weight	IG	BLINK-J	VStar overall (direct/relative)	MMVP	SAT overall (action/ego/goal/object/perspective)
0	10	none	58.67	60.73 (59.13/63.16)	72.33	44.67 (54.05/43.48/58.82/26.09/33.33)
2	1	none	57.33	59.69 (58.26/61.84)	75.00	48.67 (56.76/69.57/52.94/30.43/33.33)
2	3	none	58.00	59.69 (56.52/64.47)	73.67	52.67 (56.76/60.87/76.47/43.48/24.24)
2	10	none	52.67	60.73 (56.52/67.11)	73.33	46.67 (51.35/56.52/70.59/30.43/21.21)
4	10	none	51.33	62.30 (60.87/64.47)	70.67	45.33 (51.35/56.52/70.59/26.09/18.18)
2	10	absolute IG	46.00	59.69 (57.39/63.16)	74.33	46.00 (48.65/69.57/58.82/26.09/27.27)

All values are percentages. BLINK-J has 150 rows, VStar 191, MMVP 300, and SAT circular scoring reports 150 paired groups.

When an evaluation is queued before its adapter is committed, bind the SIQ task to the training SIQ Task ID with afterok and set WAIT_FOR_ADAPTER_DIR to the expected adapter directory. The evaluator waits for its atomic COMPLETE marker before loading weights; resumable --reuse keeps a retried bundle in the same result directory.

The same frozen evaluator runs other VLMEvalKit MCQ datasets by setting DATASET_NAME; it leaves each framework-generated user prompt unchanged. Versioned task evaluations live under experiments/output/uniar-sft/<dataset>/v1/runs/eval/. On Delta, A100 evaluation uses one model rank per GPU; A40 evaluation may set VLMEVAL_WORKERS_PER_GPU=2 and launch eight ranks on four GPUs. Shared-GPU ranks use Gloo only for barriers, while inference and result shards remain rank-local.

Canonical UniAR task evaluation sets UNIAR_ROLLOUT_DIR and atomically saves one rollouts/<dataset>/<sample-id>/ directory per generated answer. Its rollout.pt contains prompt, assistant and full token IDs plus observation and generated visual codes; metadata.json contains the unchanged framework message, decoded event stream, grids, stop/cycle metadata and workbook-facing artifact identity. RGB image decoding is deferred. A COMPLETE marker is the artifact commit boundary, and the workbook extra_records cell links back to the committed directory.

The runner fixes VLMEVAL_EVAL_ID=RUN for the lifetime of a prepared eval run; otherwise upstream's date/live-Git identity can change across SIQ retries and strand partial rank pickles in an older directory. It also sets VLMEVAL_DUMP_EVERY=1, so every completed sample updates its rank pickle. These are persistence controls only and do not modify generation, dataset.evaluate, or the upstream verifier. Canonical multi-rank inference uses VLMEVAL_SCHEDULING=dynamic: ranks claim the next pending row from one distributed Store counter, commit after every sample, and merge all prior rank shards before queue initialization. Static scheduling remains selectable for controlled continuations. When changing evaluator world size or hardware, seed a new parent-linked run with thinkmorph_sft.seed_vlmeval_resume: it copies committed rollout and completed benchmark artifacts, merges old rank shards into dataset PREV.pkl files, and rewrites workbook rollout links to the new self-contained run. When the canonical output path must remain stable, create the changed execution snapshot under <eval-run>/_continuations/<timestamp>/ and keep its OUTPUT_DIR pointed at the parent eval run. The nested snapshot and parent record preserve execution provenance while results and rollouts stay canonical.

The v2 BitCE-weight-3 long control exposes exact epoch-boundary evaluation keys for epoch 1 (step-001524), epoch 2 (step-003048), and epoch 3 (step-004572); all load the nested checkpoint adapter/ while validating the checkpoint-level COMPLETE marker. The lower-LR (1e-5) eight-epoch control exposes epoch 3 (step-004572) and step-009144 through dedicated evaluation model keys. The Grid16 lower-LR control exposes epoch 6 (step-009144) as uniar-thinkmorph-v2-r64-grid16-bitce-lambda3-lr1e-5-epoch6. The GH200-resumed Grid16 lower-LR control exposes epoch 7 (step-010668) as uniar-thinkmorph-v2-r64-grid16-lr1e-5-resume1100-epoch7. The matching 12-epoch continuation resumes its completed step-012192 state on four Delta H200s with the unchanged per-rank microbatch 4/global batch 16; four additional epochs end at step 18,288. The LLM fork preserves upstream VLMEvalKit verifier and option-prefetch logic; do not port the stricter local exact-match verifier into it. Its external judge is deepseek-v4-flash through the official OpenAI-compatible endpoint, using DEEPSEEK_API_KEY; judging defaults to non-thinking mode. API_KEY is an ignored local convenience file and must never enter Git or source snapshots.

Native models/ThinkMorph-7B task evaluation uses model key thinkmorph-7b-local. Its batch-8 adapter keeps static rank sharding, writes intermediate generated images under the eval run's generated_images/, and uses the same stable result namespace and DeepSeek verifier.

The BAGEL rank-128 ten-epoch control uses model key bagel-thinkmorph-r128-val5-10epoch. It reconstructs the audited 556-target LoRA schema on models/Bagel-7B-MoT/ema.safetensors, loads the final EMA adapter, and prefixes the unchanged benchmark question with the exact training system prompt after its input-image carrier instead of adding the native pre-image prompt. Use model-side batch one with the dynamic queue when sharing each GH200/H200 between two evaluator workers. The matched unfine-tuned baseline uses bagel-7b-mot-original: it loads the released BAGEL EMA weights with the same protocol, decoding and benchmark scoring, but does not inject an adapter. For the canonical VSP comparison use their -vsp-short model-key variants with VSP_VALIDATION_MODE=official. These use the exact ThinkMorph short custom prompt and only the final test image. The suffix affects VSP prompts only; the other five benchmarks retain their standard dataset prompts. Prompt variants use separate run/model namespaces to avoid prediction reuse.

UniAR conversion

Training corpus

The public ThinkMorph SFT corpus under data/ThinkMorph/ contains 24,990 train-only rows: Chart Refocus 6,000, Jigsaw Assembly 6,000, Spatial Navigation 6,000, and Visual Search 6,990. Each row has one problem image, two textual reasoning segments, one intermediate reasoning image, and one answer. Public training rows are not repurposed as a held-out test split.

thinkmorph_sft.prepare_cache converts every row to the canonical spatial_mmcot_trajectory_v1 and zebra_uniar_cache_v1 contracts. The exact event order is system prompt, user question and observation, then assistant thought, generated state, thought, and final answer. Rendering is owned by the same v2 ChatML path as Zebra-CoT, including exactly one newline between assistant events and the student-owned <|vision_start|> boundary. Generated images are encoded into UniAR BSQ targets; raw ThinkMorph records remain immutable.

The full-data cache is data/ThinkMorph/ThinkMorph. All its manifest rows have split train; it is unchanged by development splitting. Full-data training disables internal validation. Configuration/epoch selection instead uses the fixed data/ThinkMorph/ThinkMorph-IDVal5 development view described below. Final task evaluation uses external benchmarks, not training membership.

Fixed 16×16 generated-state cache

thinkmorph_sft.prepare_grid16_cache builds the parallel all-train cache data/ThinkMorph/ThinkMorph-Grid16 and then projects the existing IDVal5 assignments onto data/ThinkMorph/ThinkMorph-Grid16-IDVal5. It never modifies the original cache or recomputes validation membership.

For every assistant-generated target, the source RGB image is resized with preserved aspect ratio until its longest edge is 512 pixels, centered on a 512×512 RGB (128,128,128) canvas, and then encoded by the canonical UniAR BSQ encoder. The resulting patch grid is always 32×32 and its composite BSQ grid is always 16×16. User-query images are not passed through the historical 960px cap or the extra qwen-vl-utils resize stage; they go directly through the checkpoint image processor, matching the intended understanding path and its declared min_pixels=4096, max_pixels=16777216, patch size 16, and merge size 2.

OURS/config/sft/v2/thinkmorph_vanilla_r64_val5_grid16_bitlambda3_lr1e-5_8epoch.yaml is byte-for-byte equivalent in model, loss, optimizer, schedule, batching, validation, and duration to the prior lower-LR eight-epoch control, except for its run identity and Grid16 development-cache path.

UniAR v1 vanilla configuration

OURS/config/sft/v1/thinkmorph_vanilla_base.yaml is the train-only UniAR counterpart of the v1 Tetris vanilla base, not the separate BAGEL baseline. It uses all 24,990 cached rows without a validation/test split or prefix subset, one epoch, half-sample trajectory/update-uniform BitFlip in [0, 0.1], IG/EMA disabled, and the same text/bit weights, rsLoRA, optimizer, scheduler, and clip norm 5. Its current rank is 64 (Tetris's original base uses 32). Four-GPU DDP with microbatch 4 and accumulation 1 gives nominal logical batch 16. TensorBoard and JSONL log every update; update/wall checkpoints remain enabled.

evaluation.enabled=false, evaluation.at_start=false, evaluation.samples=0, and data.validation_splits=[] disable startup, periodic, resumed-step, and final internal validation. No test job is implied. The Zebra-specific probe IDs/provenance are cleared; this config targets sft train, not the Zebra-small sft inspect/probe diagnostics.

UniAR v2 constant schedule

OURS/config/sft/v2/ carries the long-horizon configuration family. Its base files retain the v1 model, data, objective, optimizer, batching, evaluation, checkpoint, and logging settings, while switching to constant LR after a fixed update-count warmup. ThinkMorph uses 50 warmup updates; the value is independent of the configured epoch count.

thinkmorph_vanilla_r64_val5_bitlambda3_16epoch.yaml is the first v2 development run: rank 64, fixed IDVal5 split, 16 epochs, Bit CE weight 3, vanilla SFT with IG/EMA disabled, logical batch 16, and clip norm 5. It saves an exact checkpoint at every epoch boundary in addition to the ordinary 100-update and 30-minute recovery points, then runs the complete 1,249-row validation exactly once per epoch. Update-periodic validation is disabled for this config. The 4/8/12/16-epoch states therefore belong to one constant-LR trajectory and can be evaluated without changing the scheduler horizon.

Fixed same-distribution development validation

data/ThinkMorph/README.md declares the two views. The development view has 23,741 training and 1,249 validation rows (approximately 5% held out): Chart Refocus 5700/300, Jigsaw Assembly 5700/300, Spatial Navigation 5700/300, and Visual Search 6641/349. Whole-group and subtype rounding takes precedence over forcing exactly 1,250 rows. This is our development split, not an author-held-out validation/test set.

The split builder owns grouping and selection:

• Preserve source cache IDs and all token/code tensors; only metadata changes. The new view uses the established absolute-shard-path zero-copy contract.
• Join original-image filename families, exact decoded RGB input and target images, and cross-task image matches into connected groups. Reconstructed jigsaw targets protect against different shuffles of the same source image.
• Natural-image resize/re-encoding checks use a conservative 128-bit dHash and thumbnail-difference rule. Sparse charts/mazes are excluded from that fuzzy rule to avoid grouping distinct diagrams merely because templates match.
• Stratify whole groups by task and available subtype: puzzle grid, maze size, and chart/search source type. A seeded whole-group subset targets 5% of rows per stratum; cross-stratum connected groups stay intact.
• Neither answer labels, model losses, nor evaluation outcomes select the membership. There are 28 multirow groups and 16 recorded near-duplicate linking edges. Source identity, decoded-pixel identity, and group leakage audits all pass. This does not prove absence of all arbitrary crops or semantically equivalent rerenders without retained source identity.

SPLIT.json records policy, source hashes and the manifest checksum; image_groups.jsonl records image/source/group/subtype provenance. The fixed manifest SHA256 is 6fc8ab060cd29efebe303c4a9c0660271d6e8f5c69c49375a466aa2f04c92c1a.

Task-specific v3 SFT uses zero-copy views of this fixed split. The Navigation view is data/ThinkMorph/ThinkMorph-Navigation-IDVal5: it retains exactly the original 5,700 train / 300 validation Spatial_Navigation rows and points to the unchanged tensor shards. Its vanilla configuration is OURS/config/sft/v3/thinkmorph_navigation_vanilla_r64_val5_lr1e-5_8epoch.yaml; only the run identity, cache view, and full-validation sample count differ from the v3 all-task control. Its controlled absolute-IG counterpart changes only the run name, enables validation IG, and applies the established all-event best setting: T→V weight/temperature 0.05/32 and V→T 10/0.25. The Navigation BitCE scale control changes only the vanilla run name and visual BitCE weight from 10 to 3; it retains eight epochs and the same fixed split. The Grid16 Navigation vanilla control instead keeps BitCE weight 10 and all training settings, but uses the zero-copy data/ThinkMorph/ThinkMorph-Grid16-Navigation-IDVal5 view. All 5,700/300 rows preserve the fixed split and use a 16×16 generated BSQ target. The completed epoch-6 (step-002196) and epoch-8 (step-002928) checkpoints are registered for VSP evaluation as uniar-thinkmorph-v3-navigation-r64-lr1e-5-epoch6 and uniar-thinkmorph-v3-navigation-r64-lr1e-5-epoch8. The matching absolute-IG epoch-6/epoch-8 VSP-short keys are uniar-thinkmorph-v3-navigation-ig-abs-best-r64-lr1e-5-epoch6-vsp-short and uniar-thinkmorph-v3-navigation-ig-abs-best-r64-lr1e-5-epoch8-vsp-short. The Grid16 navigation epoch-8 VSP-short key is uniar-thinkmorph-v3-navigation-grid16-r64-lr1e-5-epoch8-vsp-short. Reproduction (only to a new output path; existing views fail closed):

PYTHONPATH=OURS/src:experiments/src mamba run -n UniAR \
  python -m thinkmorph_sft.split_validation \
  --source-root data/ThinkMorph \
  --source-cache data/ThinkMorph/ThinkMorph \
  --output data/ThinkMorph/ThinkMorph-IDVal5 --seed 17 --workers 2

Rank-64 duration selection and full-data final training

The explicitly controlled development configs are:

• OURS/config/sft/v1/thinkmorph_vanilla_r64_val5_2epoch.yaml
• OURS/config/sft/v1/thinkmorph_vanilla_r64_val5_4epoch.yaml
• OURS/config/sft/v1/thinkmorph_vanilla_r64_val5_8epoch.yaml

These are result-producing v1 runs and use run.category: runs; their artifacts resolve below experiments/output/uniar-sft/thinkmorph/v1/runs/sft/.

The 2-epoch visual-loss scale ablation uses thinkmorph_vanilla_r64_val5_bitlambda{1,3}_2epoch.yaml. Each changes only run.name and objective.bit.weight from the rank-64 2-epoch baseline (whose bit weight is 10); text weight remains 1. These are independent runs from UniAR-SFT with their own identical 2-epoch cosine schedules and the full 1,249-row development validation.

thinkmorph_ig_abs_r64_val5_2epoch.yaml is the controlled absolute-IG counterpart of the rank-64 2-epoch vanilla configuration. It enables all-event absolute IG with T→V weight/temperature 0.05/32 and V→T 10/0.25, and records IG during the same complete 1,249-row validation. Every other model, data, BitFlip, optimizer, schedule, DDP, and logging field matches vanilla.

They differ only in run name and epoch count. All use rank 64, rsLoRA alpha 32, the existing LoRA/head learning rates 5e-5/1e-5, uniform half-sample BitFlip, IG/EMA off, global batch 16, clip norm 5, and every-step TensorBoard/JSONL. Each starts from UniAR-SFT independently and uses its own duration's cosine schedule; this is not equivalent to taking epoch-2/4 snapshots of an 8-epoch schedule. No learning-rate change is implied by creating this split.

Internal validation is enabled at startup, every 250 updates, and training completion, with all 1,249 fixed rows. It is clean teacher-forced text/visual validation, not autonomous rollout accuracy. Treat component losses and stability as supporting evidence; task-appropriate answer/visual-operation rollouts and a protected final test remain separate workflows. Do not repeatedly use a final benchmark test for epoch/rank/LR selection.

thinkmorph_vanilla_r64_base.yaml is the matching full-data template: original cache, rank 64, no held-out split and internal validation off. Set trainer.epochs to the selected duration for the final run. It preserves the same model/objective/optimizer settings as the development family. Start the final run from the base model; changing data membership is not an exact resume of a development run. Keep the held-out results from development runs, but do not call those rows held out after final full-data training includes them.

Preparation

The cache builder is multi-GPU, strict about all task counts, and resumable from committed rank manifests. A torn terminal JSONL line is discarded; an existing atomic READY marker makes a retry exit successfully.

torchrun --standalone --nproc-per-node=4 -m thinkmorph_sft.prepare_cache \
  --source-root data/ThinkMorph \
  --model-path models/UniAR-SFT \
  --output-root data/ThinkMorph/ThinkMorph \
  --attn flash_attention_2 \
  --auto-resume