Recover and verify AstroPRISM: AION encoders, dataset bridge, launcher fixes, tutorial by Sande33p · Pull Request #12 · AI-ModCon/PRISM

Sande33p · 2026-10-20T04:44:42Z

Description

AstroPRISM aligns a frozen AION-1 Base astronomy encoder with a frozen
Qwen3-4B through two small projectors (17,059,840 of 4.45 B parameters), so the
model describes calibrated galaxy images and DESI spectra in words. The
checkpoint at Sand33p/AstroPRISM
was trained with this code.

That code had never been committed. Its model card says so: the run
"contained uncommitted AstroPRISM worktree changes, so that revision alone is
not a complete source snapshot." It survived only as a git stash in one clone
of the former BaseMM_PRISM repo — reachable by no branch, removable by a stash
drop or a reflog expiry. This PR recovers it, rebased onto current main
rather than applied as-is: the stash forked 261 commits back and seven tracked
files had moved under it.

Provenance is checkable, not asserted. The uploaded config/training_config.yaml
matches src/conf/training/astroprism_joint_aion.yaml here on task,
batch_size: 2, max_steps: 200, the 576/273 token counts, and the pinned AION
revision 40541618104bab0fa85c8af68daeb867a720bb8c.

Beyond the recovery, running the result end to end found and fixed five defects
that reading the diff did not:

build_model_config dropped all 15 AION keys — the rebase replaced the
branch's inline ModelConfig with main's refactored builder, so
image_encoder_type silently fell back to "siglip" and the joint path
tried to fetch SigLIP from the Hub.
aion_cache_dir was off by one directory. huggingface_hub's
cache_dir is the hub directory, while HF_HOME has hub/ appended.
With aion_local_files_only: true there is no network fallback, so it failed
wherever the weights were staged. The released run recorded
aion_cache_dir: null, so this path was never exercised.
The published dataset could not be loaded at all — the configs read local
manifests, the HF package ships a different envelope, and the original run
used hf:// paths nothing in this repo resolves.
tools/materialize_astroprism_hf_dataset.py bridges them.
metadata/candidates.jsonl carries image captions only — 0 of 646 rows
has a spectrum caption; they live in the spectrum_caption Parquet config.
Stage 1B and the joint path had no roster at all.
Three Aurora launcher defects, independent of AstroPRISM but found with
it: --batch executed the PBS script instead of submitting it, a stale
tarball was never re-extracted, and the activate path ignored the tarball's
nested .venv-deepspeed.

Also included: docs/astroprism_tutorial.md, an Aurora walkthrough from a
clean clone to inference, and the pilot plan, roster, prompt templates and
benchmark survey under docs/applications/.

Type of Change

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that changes existing behavior)
Documentation update
Infrastructure, CI, or tooling

Compatibility

Hydra keys — all additive. Modality.SPECTRUM is new, and src/config.py
gains 15 aion_* / spectrum_encoder_type keys, each defaulted so existing
configs are unaffected. The only change to an existing value is
aion_cache_dir, in the three prism_astroprism_* configs this PR also
introduces.

Launcher flags — tools/launch_aurora_web.py gains --skip-model-staging
(opt-in). tools/launch_aurora.py changes --batch behavior: it now submits
with qsub instead of executing the generated PBS script locally. Anyone
relying on the old behavior was running their job on the login node with every
#PBS directive ignored, so this is a fix rather than a migration, but it is a
visible change.

src/model.py's _tokens_per_instance_for_modality replaces main's inline
encoder lookup because it also handles a shared encoder owner. It had to learn
main's DNA parent-modality mapping to keep BioReason working — the
resolution most worth a reviewer's eye.

How Has This Been Tested?

Everything below ran on Aurora compute nodes.

# all three stages, two-step smoke runs
Stage 1A : loss 2.9253, checkpoint saved          (job 8945735)
Stage 1B : loss 4.70,   projectors.spectrum only  (job 8945848)
Joint    : loss 6.18,   both projectors           (job 8945867)
           test/loss 5.728 (image 3.057, spectrum 14.780)

# inference against the released checkpoint         (job 8945822)
1,974 tensors loaded; 12 projector, 1,564 shared AION; route map
{image: aion, spectrum: aion}. Generated spectrum caption is
byte-identical to examples/README.md:
  "At the catalog redshift z ≈ 0.113, the spectrum shows a red, rising
   continuum and a red, rising [O III] emission line."
36 tokens, eos. The image example matches its published first sentence and
diverges only where our token budget was 320 against the release's 302.

# gates
ruff 0.16.9 over src tests tools scripts applications examples   clean
pytest -q (CI unit marker expression)      1610 passed, 1 pre-existing failure
tools/ci/check_doc_links.py                495 links resolve
tools/ci/check_site_paths.py               no new entries
launcher_contract job                      23 passed
tests/test_makefile_ci_parity.py           12 passed

Platform tested on: Aurora (12 PVC tiles, frameworks/2026.1.0,
torch 2.13.0a0+gitcf30153 XPU).
Tests skipped, and why: test_intern_s2_input_preset_requires_its_encoder_implementation
fails identically on unmodified main in this environment — it skips below
transformers<5.2.0, which CI pins, and the module here ships 5.14.1. It will
skip in CI. Full-length training was not run: these are 2-step smokes against
the release's 200 steps, so the losses show the machinery works, not learning.

New Dependencies, Data, or Assets

polymathic-aion==0.0.2 (+ jaxtyping, astropy) — the AION encoder.
Documented in the tutorial, not added to requirements/, since the
astronomy path is optional and --no-deps is required on Aurora so the
module's XPU torch is not replaced.
docs/applications/astroprism_data/overlap_candidates_646.csv.gz (39 KB)
— a derived cross-match table: identifiers, coordinates, separations and
shard/row locators. No caption text, pixels, spectra or weights. SHA-256
verified against its documentation. Sources are public and revision-pinned:
EiffL/desi_x_ls_aion
(CC BY 4.0) and
astronolan/galaxy-descriptions
(MIT), both of which permit redistributing derived work with attribution. A
NOTICE entry names them and the underlying DESI surveys.
Models pulled at runtime, not vendored: Qwen3-4B
(1cfa9a72…), AION-1 Base (40541618…).

Checklist

I have performed a self-review of my own changes
The code follows the project's style guidelines
I have commented my code where the intent is not obvious
I have made corresponding changes to the documentation
I have added tests that cover this change, or explained why none are needed
New and existing tests pass locally
No secrets, credentials, or absolute institutional paths are included

AI/LLM Assistance

AI/LLM tools were primarily used to generate code or artifacts in this PR
I have reviewed all such output line-by-line and understand it

This box is deliberately unchecked and must be ticked by a human before
merge. The recovery restores roughly 16,000 lines that were written in the
original worktree. They have been merged, tested and exercised end to end, but
nobody has read them for correctness. A reviewer who knows the AION side
should look at them, and especially at the five hand-resolved merge conflicts
listed under Compatibility.

Additional Context

Supersedes AI-ModCon/BaseMM_PRISM#179, open and unreviewed since 2026-09-22.
That PR should be closed pointing here; @n-getty owns the branch.
Dropped from the stash: 29 unrelated files it also carried (atomistic
encoders, VSP maze fixtures, DDP notes). They belong in their own PR and are
still in the stash.
The branch is large. If a reviewer would rather take it in pieces, a natural
split is: the three launcher fixes (self-contained, ~210 lines), the AION
encoders plus configs, then the data pipeline and docs.

AstroPRISM aligns a frozen AION-1 Base astronomy encoder with a frozen Qwen3-4B through two small projectors, so the model can describe calibrated galaxy images and DESI spectra. The checkpoint published at huggingface.co/Sand33p/AstroPRISM was trained with this code. That code had never been committed. Its model card says the training run "contained uncommitted AstroPRISM worktree changes", and it survived only as a git stash in one clone of the former BaseMM_PRISM repo, reachable by no branch. This commit brings it here, rebased onto current main rather than applied as-is: the stash forked at a commit 261 behind, and seven tracked files had moved under it. Provenance is checkable rather than asserted. The uploaded config/training_config.yaml matches src/conf/training/astroprism_joint_aion.yaml on task, batch_size 2, max_steps 200, the 576/273 token counts and the pinned AION revision 40541618104bab0fa85c8af68daeb867a720bb8c. Contents: - src/encoders/aion_{image,spectrum,multimodal}.py, the frozen AION encoders, plus Modality.SPECTRUM and its float32[7781] flux/ivar/mask/wavelength contract. - 11 src/data/astroprism_* modules and 9 Hydra configs for the three stages: image alignment, spectrum alignment, and the joint run. - docs/astroprism_tutorial.md, an Aurora walkthrough from environment to inference against the published weights, alongside the pilot plan, roster, prompt templates and benchmark survey under docs/applications/. - Three Aurora launcher fixes, independent of AstroPRISM but found with it: --batch executed the PBS script instead of submitting it, a stale tarball was never re-extracted, and the activate path ignored the tarball's nested .venv-deepspeed. tests/test_launch_aurora.py covers them and is wired into both launcher_contract and the make target, since its markers would otherwise leave it running in no CI job. Merge resolutions worth review. main's newer work wins where the two diverged: answer-only loss masking, aux decoder targets, generalized encoder unfreezing, the bioreason_grpo trainer dispatch, and the LoRA scopes. AstroPRISM's _tokens_per_instance_for_modality replaces main's inline encoder lookup because it also handles a shared encoder owner; it had to learn main's DNA parent-modality mapping to keep BioReason working, and it tolerates a missing ownership map so routing tests that build via __new__ still pass. main's trainer dispatch now forwards the validation and test loaders the joint path builds. The README support table gains a spectrum row and its config counts move from 26 to 29, which tests/test_docs_consistency.py enforces. Dropped from the stash: 29 unrelated files it also carried (atomistic encoders, VSP maze fixtures, DDP notes), which belong in their own PR. Verified on Aurora: 92 AstroPRISM tests, 92 multimodal contract tests, 23 launcher contracts, ruff clean, doc links resolve, and the site-path ratchet gains no entries. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…found Three defects, all found by executing the tutorial rather than reading it. 1. The published dataset could not be loaded at all. The training configs read local manifests (artifact_type astroprism_raw_images, an outputs block, a per-candidate list); the Hugging Face package ships an upload index and a bare candidates.jsonl, and the original run pointed at hf:// paths that nothing in this repo resolves. Building AstroPRISMStage1ADataset against the real snapshot failed on the manifest, so the checkpoint was not reproducible from what was published. tools/materialize_astroprism_hf_dataset.py bridges the two. It writes the four manifests the configs expect, rebuilding the per-candidate list from the band arrays, and references the payloads by absolute path so the arrays stay in the Hugging Face cache instead of being copied. With it, the real Stage 1A dataset builds: 516 train examples, image (4,160,160) float32, 425 text tokens. 2. build_model_config dropped all fifteen AION keys. They came from the inline ModelConfig the recovered branch used, which the rebase replaced with main's refactored builder, so image_encoder_type fell back to "siglip" and the joint path tried to fetch SigLIP from the Hub for an encoder the run never wanted. Verified on an Aurora node: before, the build died in ImageEncoder; after, it reaches AIONMultimodalEncoder with aion/aion and 576/273 tokens. 3. The tutorial's install step was wrong twice. huggingface-cli is gone in huggingface_hub 1.x (it is `hf download` now), and installing polymathic-aion with --no-deps leaves out jaxtyping and astropy, which the frameworks module does not supply -- aion.codecs raises ModuleNotFoundError as soon as the encoder is built. tests/test_materialize_astroprism_hf_dataset.py covers the bridge with a synthetic snapshot: no network, no gated repo, no large arrays. It asserts the artifact types the loaders demand, the absolute output paths, the per-candidate index and band order, and that a partial snapshot or a mismatched band array fails rather than producing a manifest that points at nothing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The last thing standing between the recovered code and a reproducible run was a one-level path error. All three astroprism model configs set aion_cache_dir: ${ASTROPRISM_ROOT}/hf/cache but huggingface_hub's cache_dir IS the hub directory, while HF_HOME has hub/ appended to it. Measured on an Aurora node: cache_dir=.../hf/cache -> LocalEntryNotFoundError cache_dir=.../hf/cache/hub -> resolves With aion_local_files_only: true the run cannot fall back to the network, so AION.from_pretrained failed no matter where the weights were staged. The released checkpoint's training_config.yaml records aion_cache_dir: null -- the ambient cache -- so this path was never exercised. The configs now say null and explain why. With that fixed, Stage 1A runs end to end on one node: Step 1: Loss 2.9253 Training Zone A: 100%|##########| 2/2 [00:10<00:00, 5.47s/it, loss=2.93] Checkpoint Saved. Forward 62% / backward 33% of step time, 8.98 GB checkpoint (the frozen Qwen and AION tensors are saved beside the projector that trained). That is the first confirmation that the recovered implementation, the materialized dataset and the published weights' configuration agree. The tutorial gains the AION staging step the run proved is required -- 1.8 GB at the pinned revision, which a compute node will not fetch for itself -- and section 4 now shows the smoke-run output above instead of describing it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The published metadata/candidates.jsonl carries image captions only -- 0 of its 646 rows has a spectrum caption. The spectrum captions live in the dataset's own `spectrum_caption` Parquet config instead, one shard per split (516 train / 65 validation / 65 test, all accepted), so Stage 1B and the joint path had no roster at all in the published tree. The materializer now builds one: read the three Parquet shards, join them onto the image roster by candidate_index -- which is also the raw-array index -- and write captions.jsonl beside its manifest, since the loader hardcodes that filename rather than reading the manifest's path field. The per-record review and policy blocks are derived rather than invented: caption_status and training_allowed come from the Parquet, and the split roles follow the dataset card, which permits gradients on train only. A row whose training_allowed disagrees with its split is rejected instead of quietly becoming trainable. Reaching this took five further validation layers, each hidden behind the last: the Stage 1B artifact_type, the review and training-policy blocks, candidate_indices with split counts, roster provenance that both manifests must agree on, and finally the per-record contract. With it, inference against the published checkpoint reproduces the released example exactly: At the catalog redshift z = 0.113, the spectrum shows a red, rising continuum and a red, rising [O III] emission line. 36 generated tokens, stopping on eos, byte-identical to examples/README.md. The image example matches its published first sentence and diverges only where our token budget was 320 against the release's 302. The runner's own contract check passes on the way: 1,974 tensors, 12 projector tensors, 1,564 shared AION tensors, and the joint route map {image: aion, spectrum: aion}. That is the end-to-end proof the recovery was faithful -- the restored code, the materialized dataset and the published weights agree to the token. Five tests cover the new path, including the two negative cases that matter: a missing Parquet shard and a split/training_allowed disagreement both exit non-zero rather than producing a roster that looks usable. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Running the Stage 1B and joint smokes showed the page had the pipeline shape wrong. It presented three independent launch_aurora.py invocations, but the joint config carries resume_weights_only: ${ASTROPRISM_JOINT_ADAPTER_INIT, _local/astroprism/joint_adapter_init} so the joint stage resumes from a checkpoint holding both already-aligned adapters. Launched on its own it dies with FileNotFoundError before the first step. The stage-1 checkpoints show why: 1A saves projectors.image only, 1B saves projectors.spectrum only, and tools/compose_astroprism_joint_adapters.py merges them into the init the joint run expects. That tool shipped with the recovery but was undocumented. The page now shows the dependency, documents the compose step, and records what the two new smoke runs measured: 1B at loss 4.70, joint at 6.18 with a per-modality test split (image 3.057, spectrum 14.780) that confirms the alternating batch sampler routes both modalities through one shared AION core. Two steps against the released run's 200, so the numbers show the machinery works and nothing else. The lead-in no longer overstates what was checked: all three stages ran as smoke runs, plus inference against the released weights. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Building the environment from a fresh clone of this branch showed the page starts at `python -m venv` but never says to clone or cd into the repo -- yet the next line is `pip install --no-deps -e .`, which only works from inside it. Added, with the clone on /flare rather than $HOME, matching where the caches have to live. The rest of section 1 is now verified clean-room rather than inferred: a fresh venv built exactly as written imports aion.codecs, lists all eight modalities, and resolves src through the editable install from a neutral cwd. Recorded the two facts that tell a reader their environment is right -- about 1.4 MB, because everything heavy comes from the frameworks module, and torch still 2.13.0a0+gitcf30153 with no site-packages/nvidia/, which is what a leaked CUDA wheel stack would look like. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Copilot

🟡 Changes recommended

Codec precision loss, launcher synchronization and cache defects, broken roster defaults, and a contradictory test assertion remain unresolved.

9 open issues

What changed in this PR

Recovers AstroPRISM’s astronomy workflow within PRISM, connecting AION image/spectrum encoders to Qwen3-4B through trainable projectors.

Changes:

Adds astronomy encoders, data contracts, and three-stage training presets.
Adds extraction and adapter-composition tools with regression coverage.
Updates Aurora launching, CI integration, and documentation.

File	Description
tools/prepare_astroprism_stage1a.py	Adds feature-parity preparation CLI.
tools/launch_aurora.py	Updates environment refresh and PBS submission.
tools/launch_aurora_web.py	Adds optional model-staging bypass.
tools/extract_astroprism_raw_spectra.py	Adds spectrum extraction CLI.
tools/extract_astroprism_raw_images.py	Adds image extraction CLI.
tools/compose_astroprism_joint_adapters.py	Combines trained projectors with provenance.
tests/test_zone_a_aion_training_boundary.py	Tests codec freezing and optimizer boundaries.
tests/test_prepare_astroprism_spectrum_caption_pilot.py	Tests caption payloads and selection.
tests/test_materialize_astroprism_hf_dataset.py	Tests published-dataset manifest conversion.
tests/test_launch_aurora.py	Tests environment layout and submission.
tests/test_launch_aurora_web_vla.py	Tests staging-bypass validation.
tests/test_compose_astroprism_joint_adapters.py	Tests projector composition.
tests/test_astroprism_train_entrypoint.py	Tests image-training configuration contracts.
tests/test_astroprism_stage1b_spectrum_train_entrypoint.py	Tests spectrum-training configuration contracts.
tests/test_astroprism_stage1a.py	Tests feature extraction and validation.
tests/test_astroprism_spectrum_facts.py	Tests spectral measurements.
tests/test_astroprism_raw_spectra.py	Tests spectrum extraction and integrity.
tests/test_astroprism_joint_train_entrypoint.py	Tests joint-training contracts.
tests/test_astroprism_joint_data.py	Tests joint datasets, collation, and sampling.
tests/test_aion_spectrum_model_wiring.py	Tests spectrum encoder integration.
tests/test_aion_spectrum_encoder.py	Tests spectrum encoding invariants.
tests/test_aion_model_wiring.py	Tests image encoder integration.
tests/test_aion_image_encoder.py	Tests image encoding invariants.
tests/test_accept_astroprism_stage1a.py	Tests acceptance and grouped splits.
tests/multimodal/test_modality_enum.py	Adds spectrum enum coverage.
tests/multimodal/test_make_dummy_batch.py	Tests dummy spectrum fields.
src/modules/projector.py	Adds text-scaled LayerNorm initialization.
src/modalities.py	Registers spectrum and dummy inputs.
src/encoders/aion_spectrum.py	Adds raw DESI spectrum encoding.
src/encoders/aion_image.py	Adds calibrated image encoding.
src/encoders/__init__.py	Exports AION encoders.
src/data/multimodal.py	Defers optional HDF5 import.
src/data/astroprism_stage1b_spectrum_collate.py	Adds spectrum/chat collation.
src/data/astroprism_stage1a_collate.py	Adds image/chat collation.
src/data/astroprism_spectrum_contract.py	Defines lightweight spectrum constants.
src/data/astroprism_raw_spectra.py	Extracts validated spectra atomically.
src/data/astroprism_joint_dataset.py	Unifies modality-specific datasets.
src/data/astroprism_joint_collate.py	Routes homogeneous modality batches.
src/data/astroprism_joint_batch_sampler.py	Coordinates balanced distributed batches.
src/config.py	Adds AION and projector configuration.
src/conf/training/astroprism_stage1b_spectrum.yaml	Defines spectrum alignment training.
src/conf/training/astroprism_stage1a.yaml	Defines image alignment training.
src/conf/training/astroprism_joint_aion.yaml	Defines joint continuation and evaluation.
src/conf/model/prism_astroprism_spectrum_qwen3_4b.yaml	Configures spectrum-only AION/Qwen.
src/conf/model/prism_astroprism_qwen3_4b.yaml	Configures image-only AION/Qwen.
src/conf/model/prism_astroprism_joint_qwen3_4b.yaml	Configures shared AION and separate projectors.
src/conf/data/astroprism_stage1b_spectrum.yaml	Configures reviewed spectrum data.
src/conf/data/astroprism_stage1a.yaml	Configures accepted image data.
src/conf/data/astroprism_joint_aion.yaml	Configures joint modality data.
README.md	Documents spectrum support.
Makefile	Includes generic Aurora launcher tests.
experiments/astroprism_stage1b_spectrum.yaml	Adds spectrum-training launch design.
experiments/astroprism_stage1a.yaml	Adds image-training launch design.
experiments/astroprism_joint_aion.yaml	Adds joint-training launch design.
docs/index.md	Links astronomy documentation and tutorials.
docs/applications/astroprism.md	Introduces the astronomy pilot plan.
docs/applications/astroprism_data.md	Documents roster provenance and recovery.
.github/workflows/pr_quality.yml	Runs generic Aurora launcher contracts.

_{Review effort: Balanced}

CI caught one regression and Copilot's review four more, all real. tests/test_aion_model_wiring.py asserted the pre-fix aion_cache_dir. fed2228 changed it to null -- the ambient cache, matching the released run -- because huggingface_hub's cache_dir is the hub directory itself, not its parent. I changed the configs without running the test that pinned them. It now asserts the intended behavior, with the reason. Three tools defaulted --roster to docs/astroprism/data/..., the pre-flattening path. The file is committed under docs/applications/astroprism_data/, so extract_astroprism_raw_images, extract_astroprism_raw_spectra and prepare_astroprism_stage1a all failed before opening anything when run without an explicit --roster. launch_aurora.py had a race that my own staleness fix introduced. Rank 0 re-extracts when the tarball is newer than the marker, but the other ranks waited only for the marker to exist -- so with a stale marker on disk they passed immediately and activated an environment rank 0 was deleting. The wait now tests the same condition as the unpack, with a regression check that the two stay in step. launch_aurora_web.py's --skip-model-staging rejected the AstroPRISM configs outright. load_model_group_config uses yaml.safe_load and resolve_hydra_override only selects values, so the backbone path arrives as the literal string "${oc.env:ASTROPRISM_ROOT,...}/hf/models/qwen3-4b" -- not absolute, not existing. It now expands the oc.env form before testing the filesystem. The tutorial's install list omitted wandb, which src/train.py imports and every command here sets to offline mode. Building the environment clean-room surfaced it: seven test modules failed to collect with ModuleNotFoundError. Verified in that clean-room environment: ruff 0.16.9 clean, 1592 unit tests pass (the two failures left are the pre-existing ones that also fail on main here -- a module ruff of 0.11.6 and transformers 5.14.1), 19 launcher contracts, doc links and the site-path ratchet. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

coverage_gate was the last failing check. It runs a wider selection than the unit job, which is why only it reached this: TypeError: pyarrow.parquet.core.ParquetWriter() got multiple values for keyword argument 'compression' prepare_astroprism_hf_dataset.py passed compression="zstd" to Dataset.to_parquet at two sites. datasets 5.x forwards writer kwargs as **kwargs and already sets compression itself, so ours arrived twice. Routing it through parquet_writer_kwargs does not help -- that is also forwarded verbatim and pyarrow rejects the name. The right fix is to stop overriding: datasets chooses compression per column, "none" for columns holding embedded binary and "snappy" for the rest, which is a better choice than a flat zstd across an image column. Verified in the clean-room environment: the file's 7 tests pass, and the coverage_gate selection is 1775 passed with only the two pre-existing failures that also fail on unmodified main here. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Each was verified against the code before changing it. Guard the AION codecs against blanket dtype casts. ZoneATrainer casts the whole encoder ModuleDict to the backbone dtype (trainer_zone_a.py:473), which reaches both AION encoders through nn.Module._apply, so the .float() they call at construction never sees it. A bf16 round-trip restores the dtype but not the values -- measured max delta 9.5e-4 on a float32 weight -- and these codecs quantize hard, so rounded weights can change token ids. That hazard is latent rather than active, and the distinction matters. _configure_encoder_for_training (line 491) runs after the blanket cast and skips encoder.to(dtype=...) for any encoder exposing set_shared_encoder_trainable, which both AION encoders do. The released checkpoint proves the protection held: its backbone is BF16, so the cast did fire, yet all 873 saved codec tensors are float32 and none equals its own bf16 round-trip -- they were never rounded. Inference against those weights is unaffected by this commit, and a retrain should tokenize identically. What the _apply override buys is independence from that ordering. Any future path that touches the encoder between the blanket cast and the protective pass, or any new call site that casts without pairing it, would silently round the codec today. The new tests assert values rather than dtypes -- coverage the suite did not have -- and both fail against the unpatched encoders. --skip-model-staging rejected the AstroPRISM configs outright. load_model_group_config uses yaml.safe_load and resolve_hydra_override only selects values, so the backbone arrives as the literal string "${oc.env:ASTROPRISM_ROOT,...}/hf/models/qwen3-4b" -- neither absolute nor existing. It now expands the oc.env form before testing paths. Skip mode also left two traps. It touched the persistent model_ready marker without copying anything, so a later staging-enabled launch on the same node would see it, skip its copy loop and fail offline against an empty cache; the marker is no longer published, and the wait on it is bypassed in skip mode. And the HF cache exports still named /tmp/huggingface, which in skip mode is empty, so AION (aion_cache_dir: null) searched it and failed offline even with the weights in the shared cache. Both export sites now follow `use_shared_venv or skip_model_staging`; the third use_shared_venv branch selects the Python environment and is deliberately untouched. Verified: a generated skip-mode script parses, validation passes, no marker wait, and HF_HOME no longer points at the empty node-local cache. 1777 tests pass, ruff 0.16.9 clean. The two failures left are the pre-existing ones that also fail on unmodified main in this environment. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Sande33p and others added 6 commits October 19, 2026 05:33

Sande33p requested a review from a team as a code owner October 20, 2026 04:44

Sande33p requested a balanced review from Copilot October 20, 2026 04:47

Copilot started reviewing on behalf of Sande33p October 20, 2026 04:47

View session

Copilot AI reviewed Oct 20, 2026

View reviewed changes

Sande33p and others added 3 commits October 20, 2026 06:21

Sande33p force-pushed the astroprism-implementation branch from 714c513 to 4569534 Compare October 20, 2026 07:40

Sande33p mentioned this pull request Oct 20, 2026

Open-source readiness: PRISM v0.1.0 #2

Open

20 of 31 tasks

Navigation Menu

Repository navigation

Recover and verify AstroPRISM: AION encoders, dataset bridge, launcher fixes, tutorial - #12

Recover and verify AstroPRISM: AION encoders, dataset bridge, launcher fixes, tutorial#12
Sande33p wants to merge 9 commits into
mainfrom
astroprism-implementation

Sande33p commented Oct 20, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Sande33p commented Oct 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of Change

Compatibility

How Has This Been Tested?

New Dependencies, Data, or Assets

Checklist

AI/LLM Assistance

Additional Context

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Sande33p commented Oct 20, 2026 •

edited

Loading