Fix Bagpiper public-checkpoint training, resume, and inference - #6822

Codecov Report

❌ Patch coverage is 85.22727% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.67%. Comparing base (158e495) to head (36dae7f).
⚠️ Report is 2 commits behind head on master.

@@            Coverage Diff             @@
##           master    #6822      +/-   ##
==========================================
+ Coverage   73.48%   73.67%   +0.18%     
==========================================
  Files         860      862       +2     
  Lines       80440    80629     +189     
==========================================
+ Hits        59112    59403     +291     
+ Misses      21328    21226     -102     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

Merge Risk: 🟡 Moderate · up to 91969

Existing checkpoints trained with a different vocabulary order can silently produce incorrect outputs. Provide a legacy compatibility or migration path before merging, while retaining deterministic ordering for public checkpoints.

Security Architecture Review

Security architecture risk: 🔵 Low · up to 91969

The new loading paths retain restrictive deserialization defaults. The main design risk is silent incompatibility with older saved weights using a different token order. External-file trust and interrupted recovery remain partly unverified.

Retained concerns

• Medium · architecture · inferred: Text-first, name-sorted vocabulary construction changes embedding-row meaning for checkpoints trained under a different pre-change insertion order. Native and automatic DCP loading check tensors but do not compare the original token layout, so equal-shaped legacy weights can be silently reinterpreted after upgrading. This also complicates rollback to order-dependent code. The documented released layout is supported; whether affected legacy checkpoint/configuration pairs exist is unverified.

• Fix all pre-merge checks with AI

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Autopilot is currently an internal CodeRabbit preview.

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

_{Comment @coderabbitai help to get the list of available commands.}

Thanks @sw005320 for the review and for rerunning the failed jobs. The follow-up changes are pushed:

1. Documentation: Removed both EXPERIMENT_LOG.md and REPRODUCE.md, including the personal paths and remote instructions. Usage and reproduction commands are now in the two recipe READMEs, and the validation summary is in the PR description.

2. Vocabulary compatibility: I reverted the canonical-order change entirely, keeping the existing YAML mapping-order behavior. speechlm_job.py and its vocabulary tests match master, so this PR introduces no token-order change or checkpoint migration. Both recipe READMEs document preserving the original checkpoint configuration's mapping order and using yaml.safe_dump(..., sort_keys=False). This decision is also recorded in the PR description.

3. Coordination with Bagpiper on the command line and the MCP server #6810: Adopted your if not isinstance(audio_features, torch.Tensor) form and unwrap .last_hidden_state only for structured outputs, without passing return_dict. The regression tests cover Qwen2.5/Qwen3 with both output forms. I also ran the two exact proposed Bagpiper on the command line and the MCP server #6810 encode_batch tests alongside the SpeechLM suite: 404 tests passed. This follows your proposed merge/rebase plan for Bagpiper on the command line and the MCP server #6810.

All four CI gates passed on 1b2ae0934a. The subsequent master merge is 59f6d653cb, which leaves these reviewed files unchanged and triggers a fresh CI run.

Could you please re-review the updated PR when you have a chance?

@jctian98 — a gentle nudge on this one, and an offer.

As of my last pass everything is cleared except the recipe point: saying at the top of both run.sh files that they launch training only, and giving them the export and inference steps so ./run.sh goes from public weights to decoded output like the rest of egs2.

The nudge is only because #6810 is queued behind this. We agreed this PR merges first with the isinstance form, and that I then rebase #6810 and drop my audio.py hunk — which I am ready to do the moment this lands. While both sit open, the conflict in encode_batch and test_audio.py keeps coming back on every merge from master.

If the recipe work is awkward to fit in this week, I am happy to push it to your branch (maintainer_can_modify is on): the header comment on both run.sh files, plus a --stage train|export|infer switch that runs export_checkpoint on the latest complete DCP and then espnet2.speechlm.bin.inference with the released inference_audio.yaml / inference_text.yaml. You would keep the review; I would just be typing. Say the word and I will open it as a commit on your branch for you to look over.

@jctian98 — I have pushed the recipe change to your branch rather than leave it blocking (04779a9); revert or rewrite it freely, it is your PR.

What it does. egs2/TEMPLATE/speechlm1/run.sh is a staged entry point that both recipes' run.sh now call:

Nothing that worked before changes. The staged runner consumes only the options it owns and forwards everything else to train.sh verbatim — parse_options.sh could not do that, since it exits on the first option it does not know, which is every training option. So ./run.sh --ngpu 8 --stats-dir ... --train-unregistered-specifier ... behaves exactly as it did.

Defaults. --stage and --stop-stage both default to train, so ./run.sh still trains and stops: these recipes prepare no data, and a decode needs a test manifest from you. Naming one stage runs that stage alone. I took that as the --stop-stage escape in the review rather than making a decode the default — say the word if you would rather it ran through.

The part I think is the real win: ./run.sh --stage infer --export-path /path/to/bagpiper-sft/model.pt --inference-config .../inference_audio.yaml --test-unregistered-specifier ... goes from the published weights to decoded output with no training and no export step. Both recipe READMEs now show that, and the export sections show ./run.sh --stage export with the raw module command kept for naming an older checkpoint.

Checked: bash -n and shellcheck -S style clean on all three scripts; the stage dispatch exercised by hand, including that the export stage picks step_20 over step_9 (numeric, not lexicographic) and skips a step_30 with no .metadata. Not run end to end — that needs the GPUs and the manifests.

With this in, I believe the only thing left on this PR is your call on the vocabulary-order question; #6810 then rebases and drops its audio.py hunk as agreed.

One more on your branch (3b32e2a), on the same theme as the last one: the recipe now carries the training schedule, not just the launcher.

The three stages were in the README as commands to retype, each carrying its config, its output directory and the step number of the checkpoint to resume from:

--train-config conf/tuning/train_pretrain.yaml --output-dir exp/pretrain
--resume-path exp/warmup/checkpoints/step_10000

They are now in run.sh:

./run.sh --train-stage warmup|pretrain|sft --ngpu 8 --stats-dir ... --train-unregistered-specifier ...

Two things that fall out of it, and both are the reason I think this belongs in the recipe rather than the prose:

The step numbers are gone. A stage starts from the latest complete checkpoint of the stage before it — the highest step_* that has a .metadata — so step_10000 and step_600000 no longer need to be right in the README, and the note telling readers to update them when they change a schedule is gone with them.

Re-running an interrupted stage continues it. Once a stage's own output directory has a checkpoint, the chaining is skipped and --resume-path is not passed, which is what train.sh needs to restore the optimizer and step rather than start the stage again with a fresh one. That was the distinction the README drew in prose and the command line could get wrong in either direction.

An explicit --resume-path still wins over both, which is how Bagpiper-TTS starts from the published base.pt; its run.sh says so at the top, since it has a single stage.

Checked by hand with a stubbed train.sh: each stage forwards the right config and output directory; pretrain passes --resume-path exp/warmup/checkpoints/step_10000 when that exists and nothing when exp/pretrain already has its own; an unknown --train-stage is named. bash -n and shellcheck -S style clean on all three scripts. The recipe has no data, so nothing was trained.

Reworked the stages on your branch (a5fd550) after review here: they are one ordered sequence now, not two switches.

Before, --train-stage chose between three training configurations and --stage chose between train, export and infer — two axes over what is really one pipeline. A recipe now declares its stages and the shared runner walks them in order:

--train-stages 'warmup:conf/train.yaml:exp/warmup
                pretrain:conf/tuning/train_pretrain.yaml:exp/pretrain
                sft:conf/tuning/train_sft.yaml:exp/sft'

./run.sh --ngpu 8 ...                      # warmup → pretrain → sft → export
./run.sh --stage sft --stop-stage sft ...  # one of them
./run.sh --stage infer ...                 # decode, skipping training entirely

--stage and --stop-stage take a stage name and select a range of that one sequence. Bagpiper-TTS declares a single training stage and gets the same pipeline.

The default last stage is export, not infer: the data to decode is not prepared in these recipes, and a bare ./run.sh should not end in an error about a manifest after days of training. Asking for a later stage carries the default along, so --stage infer does decode — and with --export-path and --train-config pointing into a downloaded model directory, it decodes the published weights without training anything, which is the route the READMEs now show first.

Two details worth flagging:

• Chaining beats step numbers. A stage starts from the latest complete checkpoint of the one before it, so step_10000 and step_600000 are gone from the README along with the note telling readers to update them when a schedule changes.
• bash 3.2. ${array[-1]} is a bash 4.2 feature and macOS still ships 3.2, where it fails with "bad array subscript". Found it running the recipe here; the runner uses ${array[${#array[@]} - 1]}.

Rehearsed with a stand-in train.sh that leaves a checkpoint behind: the full sequence chains warmup → pretrain → sft and exports from SFT's latest; --stage sft runs SFT and the export; --stage infer with published weights decodes; a stage whose predecessor has no checkpoint, a missing --inference-config, a missing test specifier and an unknown stage name each stop with their own message. bash -n and shellcheck -S style clean on all three scripts. No GPU, so nothing was actually trained or decoded.

Moved the stages into the recipes (36dae7f), after review here: run.sh should say what the recipe does, not delegate it.

./run.sh                           # warmup, pretraining, SFT, export, inference
./run.sh --stage 3 --stop-stage 3  # SFT alone
./run.sh --stage 5 ...             # decode, training nothing

Each stage is a block in run.sh carrying its own configuration, output directory, data and starting point, in the usual egs2 shape — numbered --stage / --stop-stage, kaldi option parsing, a log line per stage. Bagpiper has five; Bagpiper-TTS has three (SFT, export, inference). The shared runner I added earlier is gone.

egs2/TEMPLATE/speechlm1/stage_utils.sh keeps the two things worth not repeating in both recipes: finding the latest complete step_* DCP, and the rule that a training stage continues its own output directory when it has checkpoints and otherwise starts from the stage before it.

One bug worth naming, because it is the kind that reads as working: log() wrote to stdout, and resume_argument logs while its stdout is captured by the caller — so the timestamp line was being passed to train.sh as an argument:

train.sh: ... --python python 2026-10-02T14:20:52 starting from exp/warmup/checkpoints/step_100 --resume-path exp/warmup/...

log writes to stderr now. Found by running the recipe with a stand-in train.sh, not by reading it.

Rehearsed the same way as before: the five stages chain warmup → pretrain → SFT and export from SFT's latest; --stage 3 --stop-stage 3 runs one; stage 5 decodes given weights and a manifest; and the refusals each have their own message — no --resume-path on a fresh TTS run, missing data, missing --inference-config, missing test specifier, no complete checkpoint to start from. bash -n clean; shellcheck -S style clean apart from the two findings train.sh already has, which are the parse_options.sh idiom.