exporter: fall back to earlier step on broken final checkpoint - #242
Merged
Merged
Conversation
tunix PeftTrainer can leave the highest-numbered orbax checkpoint with an incomplete manifest if the trainer exits right after save_checkpoint is queued at end-of-training. Earlier saves taken DURING training are structurally complete since they had time to flush. On FileNotFoundError loading the newest step, fall back to next-highest and retry. seed7/99/101 step 300 saw exactly this; step 250 is fine.
PeterLi-jpg
added a commit
that referenced
this pull request
May 3, 2026
tunix PeftTrainer saves composite checkpoints with model_params/ and optimizer_state/ as siblings under each step dir. PyTreeCheckpointer.restore needs the inner pytree dir (model_params), not the composite root. Passing the composite root produces "No structure could be identified" because manifest.ocdbt lives inside model_params/, not at the top level. Fix: probe for model_params/ subdir; restore from it if present, else fall back to the leaf (legacy flat layout). Combined with PR #242 step fallback, this fixes Stage 3-export for seed7/13/42/99/101.
PeterLi-jpg
added a commit
that referenced
this pull request
May 3, 2026
) tunix PeftTrainer saves composite checkpoints with model_params/ and optimizer_state/ as siblings under each step dir. PyTreeCheckpointer.restore needs the inner pytree dir (model_params), not the composite root. Passing the composite root produces "No structure could be identified" because manifest.ocdbt lives inside model_params/, not at the top level. Fix: probe for model_params/ subdir; restore from it if present, else fall back to the leaf (legacy flat layout). Combined with PR #242 step fallback, this fixes Stage 3-export for seed7/13/42/99/101.
PeterLi-jpg
added a commit
that referenced
this pull request
May 3, 2026
* exporter: load model_params subdir for composite orbax checkpoints tunix PeftTrainer saves composite checkpoints with model_params/ and optimizer_state/ as siblings under each step dir. PyTreeCheckpointer.restore needs the inner pytree dir (model_params), not the composite root. Passing the composite root produces "No structure could be identified" because manifest.ocdbt lives inside model_params/, not at the top level. Fix: probe for model_params/ subdir; restore from it if present, else fall back to the leaf (legacy flat layout). Combined with PR #242 step fallback, this fixes Stage 3-export for seed7/13/42/99/101. * exporter: handle bare "layers" + digit segment layout tunix orbax save writes the layer index as a SEPARATE path segment after a bare "layers" key: ("layers", "0", "attn", "q_einsum", "w_lora_a", "value") The previous _LAYER_RE only matched compound segments like "layer_0" or "layer.0", so it never recognized the tunix layout and the export raised "no qwix LoRA leaves found". Verified locally on seed42 production checkpoint: 248 LoRA leaves at layers/<i>/attn/{q,kv}_einsum/w_lora_{a,b}/value paths.
PeterLi-jpg
added a commit
that referenced
this pull request
May 4, 2026
tunix PeftTrainer can leave the highest-numbered orbax checkpoint with an incomplete manifest if the trainer exits right after save_checkpoint is queued at end-of-training. Earlier saves taken DURING training are structurally complete since they had time to flush. On FileNotFoundError loading the newest step, fall back to next-highest and retry. seed7/99/101 step 300 saw exactly this; step 250 is fine.
PeterLi-jpg
added a commit
that referenced
this pull request
May 4, 2026
) tunix PeftTrainer saves composite checkpoints with model_params/ and optimizer_state/ as siblings under each step dir. PyTreeCheckpointer.restore needs the inner pytree dir (model_params), not the composite root. Passing the composite root produces "No structure could be identified" because manifest.ocdbt lives inside model_params/, not at the top level. Fix: probe for model_params/ subdir; restore from it if present, else fall back to the leaf (legacy flat layout). Combined with PR #242 step fallback, this fixes Stage 3-export for seed7/13/42/99/101.
PeterLi-jpg
added a commit
that referenced
this pull request
May 4, 2026
* exporter: load model_params subdir for composite orbax checkpoints tunix PeftTrainer saves composite checkpoints with model_params/ and optimizer_state/ as siblings under each step dir. PyTreeCheckpointer.restore needs the inner pytree dir (model_params), not the composite root. Passing the composite root produces "No structure could be identified" because manifest.ocdbt lives inside model_params/, not at the top level. Fix: probe for model_params/ subdir; restore from it if present, else fall back to the leaf (legacy flat layout). Combined with PR #242 step fallback, this fixes Stage 3-export for seed7/13/42/99/101. * exporter: handle bare "layers" + digit segment layout tunix orbax save writes the layer index as a SEPARATE path segment after a bare "layers" key: ("layers", "0", "attn", "q_einsum", "w_lora_a", "value") The previous _LAYER_RE only matched compound segments like "layer_0" or "layer.0", so it never recognized the tunix layout and the export raised "no qwix LoRA leaves found". Verified locally on seed42 production checkpoint: 248 LoRA leaves at layers/<i>/attn/{q,kv}_einsum/w_lora_{a,b}/value paths.
5 tasks
PeterLi-jpg
added a commit
that referenced
this pull request
May 4, 2026
orbax CheckpointManager.save() is async, so the last-step save can still be in flight when the trainer process exits. The result is a final checkpoint dir (e.g. orbax/300/) with data files but no manifest, which downstream loaders reject with 'No structure could be identified for the checkpoint at .../orbax/300'. Witnessed on seed7/seed99/seed101 step-300 ckpts; the exporter step-fallback (PR #242) is a workaround, this is the fix. Reach through tunix's CheckpointManager wrapper to the underlying orbax manager and call wait_until_finished() after trainer.train() returns. Wrapped in try/except so a missing manager (eval-only paths) does not crash the script.
5 tasks
PeterLi-jpg
added a commit
that referenced
this pull request
May 4, 2026
…o-delete (#248) bundle of 5 fixes for tpu/launch_5seeds_tunix.sh hit during tonight's pipeline run: a) repo-clone check now tests ~/bohdi-lora/.git, not just the dir, so a pre-populated ~/bohdi-lora/checkpoints/... no longer skips clone and trips the next git fetch. aborts loudly if the dir exists without .git/ instead of silently clobbering. b) shell-side watchdog around each eval invocation. kills the eval with rc=124 after VLLM_WAIT_TIMEOUT (default 600s) without a Running vllm-tpu container, or rc=1 the moment a container is observed exited (logs ExitCode). overrides the killed-process signal exit so eval_fail_count tracking sees the right rc. c) cleanup_eval restarts the tpu-runtime container so chip state is reset between Stage-4 configs. previously the next config's vllm inherited dirty device state and failed to init. d) Stage 3 train invocation now skips when checkpoints/seed_<N>/orbax/ 250/_CHECKPOINT_METADATA exists locally (resumed from GCS). step 250 is the last reliably-complete save (300 is invalid per #242). e) cleanup() trap's delete_vm gated behind LAUNCHER_DELETE_VM_ON_EXIT, default off. successful runs no longer auto-destroy spot slots that we cannot get back due to TRC capacity contention. PREEMPTED retry path keeps deleting because preempted is terminal.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
tunix PeftTrainer can leave the newest orbax checkpoint structurally invalid if the trainer exits right after save_checkpoint() is queued at end-of-training. seed7/99/101 hit this on step 300; step 250 loads fine.
Iterates checkpoint candidates newest-first; on FileNotFoundError falls back to next-highest. Prevents the same export failure for fresh training runs (seed13/42).