[ramanathan831]

Summary

• Make NVIDIA-TAO/tao-automl's packages/automl-core source installation the documented AutoML setup, using main by default with an explicit named-development-branch override. A published core wheel and a persistent TAO checkout are not prerequisites.
• Make skill-driven setup ask whether to use main or a user-supplied development branch. Preserve resolved commits/build identities for existing experiments; never silently fall back to another branch or upgrade on resume.
• Match Framework's optional dependency to the source distribution metadata, nvidia-automl-core==0.1.0, without changing either repository's package version. Explain why the separately published nvidia-tao-automl application is not the same distribution.
• Declare the existing gepa==0.1.2 requirement directly in Framework's prompt extra, update the lockfile, and retain explicit post-sync installation / uv run --no-sync so ordinary Framework setup does not fetch AutoML source.
• Update the missing-engine diagnostic, setup/algorithm/SLURM/prompt guides, and both skill-discovery paths.

Validation

• Fresh remote source install from codex/repository-neutral-automl at bd21cbd3a6aa24e3277c0d232dab1614f935c03e: confirmed nvidia-automl-core==0.1.0, recorded Git provenance, and compatible core/GEPA dependencies.
• Built the Framework wheel and checked its AutoML dependency metadata. With the wheel and source-installed core, both SFT and GEPA CLI help commands and the packaged SFT structural plan passed outside both source checkouts.
• Existing CPU adapter suite: python -m pytest -o addopts='' --confcutdir=tests/automl tests/automl — 14 passed.
• Invalid branch rejected without fallback or modification of the previously installed core; missing-engine diagnostic checked in an empty environment.
• Skill validation, matching discovery paths, and 127 local Markdown links/anchors passed.
• Ruff check/format, project pre-commit hooks, base pre-commit hooks (including secret scanning), lockfile validation, and git diff --check passed.

Scope

This is package/setup/skill validation, not GPU end-to-end validation. No training jobs, paid LLM requests, or full CUDA-environment installation were run for this change. No new unit-test suite or package release is introduced. Optional CPU-startup observer assets and backward-compatible controller-template guards are included; this PR does not change the core search/recovery-ledger implementation.

Review follow-up (2832a95)

• Preserve the existing direct-CLI SLURM launch; cross-node startup probing is opt-in through AUTOML_STARTUP_HELPER. An opted-in failure never falls back silently.
• Diagnose invalid/exhausted CPU allocation counts before helpers or traps; reject unapproved action names outside run/resume; handle bootstrap signals/failures without leaking EX_TEMPFAIL into site automatic-requeue policies.
• State the compatible core version, source-access requirement, reviewed-wheel alternative, and immutable commit/build provenance explicitly.
• Validation: 46 local shell/process scenarios; 14 existing CPU adapter tests with core 0.1.0; both CLI help paths, isolated missing-engine diagnostic, and structural plan; both full GitHub pre-commit configurations; Ruff, frozen lock consistency, skill validation/discovery, and 50 local Markdown links/anchors.
• These follow-up checks launched no SLURM allocations, GPU training, or paid LLM calls. Native scheduler requeue and cross-node storage behavior were not revalidated by the local process fixtures.

[lfengad]

A few concerns about backward compatibility of the existing AutoML controller launch path. Details inline.

[lfengad]

New conflicts might be resolved. Thanks!

[lfengad]

Thanks for the PR. Two things I'd like confirmed before merge.

1. nvidia-automl-core 1.0.0.dev0 → 0.1.0 (looks like a version downgrade)

• Any environment that already has the 1.0.0.dev0 local wheel installed will conflict with the new [automl] / [automl-prompts] pin (==0.1.0) and need a reinstall. Is that intended, and is 1.0.0.dev0 formally retired?
• cosmos_framework/automl/workload.py and prompts.py import automl_core.contracts, recovery, runner, prompt_runner. Compatibility with 0.1.0 is currently evidenced only by the 14 CPU adapter tests; no SFT/GEPA run on GPU was done. Could we get at least one small end-to-end SFT run (or an explicit sign-off from the core owners that 0.1.0 is API-compatible with 1.0.0.dev0)?
• Source install tracks tao-automl main by default and is not pinned to a commit, so two installs at different times can yield different cores under the same ==0.1.0. Should the docs recommend pinning a commit (or tag) rather than only recording it afterwards?

2. examples/automl/slurm_controller.sbatch behaviour changes

This template was only added in #288 (2026-10-17), so there is little prior usage, but anyone who copied it already would see differences:

• New early validation exits with code 2 before the CLI runs: AUTOML_CONTROLLER_MAX_ALLOCATIONS must now match ^[1-9][0-9]*$ (previously whatever the CLI accepted, e.g. a zero-padded value), and SLURM_RESTART_COUNT must be a decimal integer.
• A new allocation-cap check (generation >= cap → exit 2) now runs in the script, in addition to whatever the core CLI enforces. Please confirm the two agree on semantics (0-based generation vs. count) so the script never rejects a launch the CLI would accept, and vice versa.
• AUTOML_ACTION defaults to run, so the existing launch is unchanged. Good. The startup helper is opt-in via AUTOML_STARTUP_HELPER; please confirm that with it unset the final exec command line is byte-for-byte the previous one.
• slurm_controller_watch.py (+182 lines) only has local shell/process scenario checks and no test in the repo. Could a small unit test for registry/tick be added, or could this move out of tools/ until it has coverage?

[ramanathan831]

Addressed in 1b1c1bd. Replying point by point to the new review:

1. Core version and reproducibility

• Existing 1.0.0.dev0 environments: yes, that wheel does not satisfy the new ==0.1.0 extras. The pin follows the core source package's declared version; it is not a claim that arbitrary development builds are interchangeable. The transition guidance now explicitly requires a separate environment for a new experiment. Keep old environments and exact artifacts for existing runs; do not downgrade a live controller, relabel wheels, or rewrite ledgers. The former pin is no longer the supported new-install requirement in these extras, but this PR does not establish formal publisher retirement of every historical artifact. Cold migration remains limited to what its existing protocol validates, not arbitrary version changes or GEPA migration.
• GPU evidence: there is an existing 2026-10-19 H100 SFT verification, beyond the 14 CPU tests. I rechecked its saved native receipts: two random-search candidates, 64 WTS training records, five LoRA steps each on eight H100s, DCP-to-HF export, and 16 held-out evaluation questions each. Both completed and scored 8/16. Both controller builds declared core 0.1.0. This was a pass after explicit controller recovery from a shared-drive I/O error, with no repeated training or extended budget/deadline. The docs now identify the initial source commits and retain the corrected engine fingerprint. This is historical, build-specific SFT evidence—not an uninterrupted run, a GPU rerun of this latest PR head, or a new GEPA qualification. The four named core modules also import successfully in the tested 0.1.0 environment.
• Moving branches: agreed. The docs now recommend resolving the selected branch to a full commit before installation for reproducible runs/deployments, with an executable example and post-install provenance verification. main remains the exploratory default; reproduction uses the recorded commit or exact wheel, not a moving branch.

2. Batch compatibility and observer coverage

• Early validation: fixed the unnecessary zero-padding incompatibility. Caps such as 003 and restart counts such as 0002 now work. Comparison uses normalized decimal copies without octal parsing or integer overflow; the original cap argument is passed unchanged to the CLI. Malformed inputs still fail before startup.
• Allocation semantics: both guards count the initial allocation. Cap 3 permits generations 0, 1, and 2; generation >= 3 is exhausted. The guide distinguishes the shell's environment precheck from the core's authoritative native scontrol checks. They agree on the numerical boundary, but shell acceptance cannot guarantee acceptance of ownership, resources, ledger continuity, or deadlines by the core.
• Original launch argv: confirmed for valid default launches with AUTOML_STARTUP_HELPER and AUTOML_ACTION unset. Checked-in batch tests assert the original SFT/prompts argument ordering and values, quoting, and child PID equality proving direct exec rather than a wrapper process.
• Observer tests: added registry, readiness/locking, registration, and tick tests. They cover owner/identity conflicts, generation boundaries, non-running jobs, scheduler/storage errors, and sanitized failure receipts. The new AutoML helpers CI workflow runs the helper suite on a CPU runner without the training stack.

Validation: 172 helper tests (including 69 new cases) passed using the CI command in a fresh Python 3.12 venv with only pytest/jsonschema and their dependencies; 14 adapter tests and 28 installed-core tests also passed. Changed-file pre-commit, Ruff, lock consistency, shell syntax, and documentation-link checks passed. The new observer/batch tests use mocked scheduler/controller children; no new GPU/SLURM jobs were launched for this revision.

[lfengad]

Related follow-up to #277 / remaining coverage from #155: automatic HF export for LoRA runs.

I checked the latest head (113bdc4b). HFExportCallback._gather_weights (hf_export.py:226) still exports parameters individually, without merging the LoRA update into the base weights or removing adapter-only keys. A standard HF model loading this export would not reconstruct the trained LoRA update; it may ignore the unexpected adapter keys and effectively use the unmerged base weights.

This matters when LoRA is enabled together with this callback—for example, the VideoPhy Edge experiment already enables automatic HF export (videophy2_sft_edge.py:204). The separate offline DCP-to-HF exporter already merges and strips adapters (export_vlm_dcp.py:183), so this is specifically a gap in the automatic callback path, not a claim that LoRA training or offline export is generally broken.

Could you please cover this as well, either here if appropriate or in a small follow-up PR? #155 contains a proposed callback-side implementation that could be adapted to the current LoRA utilities rather than duplicating the core LoRA support.

Suggested coverage:

• Merge each adapter into the exported base weight without modifying the live training weights, and omit adapter-only keys from the plain HF export.
• Verify that a reloaded HF export reproduces the adapted model's eval outputs within dtype tolerance.
• Cover FSDP all-gather ordering across ranks and preserve the existing non-LoRA export behavior.

This is a code-path finding, not a new GPU reproduction. It is inherited from the current base, not introduced by this source-install/setup PR; a separately tracked fix would be fine.

[lfengad]

Hi @ramanathan831, while syncing this repo with i4 we found a small issue in LossSpikeRollback (added in #277). It isn't part of this PR, but could you please fix it, here or in a follow-up?

The learning-rate backoff takes effect one update late when a scheduler is present.

In ImaginaireTrainer, _optimizer_step() runs optimizer.step() and then scheduler.step() (trainer/init.py#L776-L777, L794) before callbacks.on_before_zero_grad(), which is where the guard rolls back and calls _scale_lr.

In the scheduler branch of _scale_lr, only sched.base_lrs is rescaled. The scheduler has already set the param-group lr for the next update, so that update still runs at the old rate. The backoff only applies from the update after that, which delays the protection right when the run is unstable.

The existing tests don't catch it because _step() in loss_spike_rollback_test.py calls on_before_zero_grad before scheduler.step(), which is the opposite of the trainer's order.

Suggested fix (we verified it in i4): also rescale the current param-group rates in that branch.

        schedulers = _real_schedulers(scheduler)
        previous_scale = self._lr_scale
        ...
        if schedulers:
            for sched, bases in zip(schedulers, self._base_lrs):
                sched.base_lrs = [base * self._lr_scale for base in bases]
            # The scheduler already stepped for this update, so its next recompute comes
            # after it. Rescale the current rates too, so the backoff applies right away.
            for opt in _real_optimizers(optimizer):
                for group in opt.param_groups:
                    group["lr"] = group["lr"] * self._lr_scale / previous_scale
            return

And a regression test that follows the trainer's order. It fails without the fix:

def test_lr_backoff_applies_to_the_next_update():
    model, optimizer, scheduler, callback = _harness()

    def trainer_step(grad_scale):  # ImaginaireTrainer order
        for parameter in model.parameters():
            parameter.grad = torch.full_like(parameter, grad_scale)
        callback.on_before_optimizer_step(model, optimizer, scheduler, None)
        optimizer.step()
        scheduler.step()
        callback.on_before_zero_grad(model, optimizer, scheduler)
        optimizer.zero_grad(set_to_none=True)

    for _ in range(20):
        trainer_step(0.01)
    trainer_step(10.0)

    assert callback.rollbacks == 1
    # The rate the next optimizer.step() uses, before any further scheduler.step().
    assert optimizer.param_groups[0]["lr"] == pytest.approx(0.01 * 0.8 * 0.5)

The callback is disabled by default, so only runs that enable it are affected. Thanks!

[ramanathan831]

Confirmed the delayed-backoff report, and implemented the fix in #307 (45be0ce).

The guard now updates both scheduler base rates and the current optimizer-group rates by new_scale / previous_scale. This preserves the current schedule multiplier while applying protection to the very next optimizer update. Recovery and repeated backoffs use the same ratio; an unchanged zero scale avoids division by zero.

Both existing step helpers now follow the actual trainer order. Before changing the implementation, three new cases failed: constant schedule, decaying schedule, and multiple optimizers/groups. After the fix, all 17 rollback tests pass, including an optimizer pre-step hook that records the rate actually consumed by the next update. The tests also cover recovery and the no-scheduler path, and now join the existing L0 CI selection.

Changed-file Ruff, both pre-commit configurations, and Pyrefly passed. These are optimizer/callback regression tests, not a new full training run. The callback remains disabled by default. As suggested, this is a separate follow-up against main because the bug is inherited from #277 rather than introduced by #291.

[ramanathan831]

Confirmed the automatic LoRA export report, and implemented the fix in #308 (7a54b49). The separate offline exporter already merges adapters; the gap was the automatic callback.

The callback-side approach from #155 is adapted to the current LoraInjectedLinear, including RSLoRA and mixed precision. A shared non-mutating merge helper is also used by the existing in-place/offline merge. The callback gathers base/A/B tensors in a deterministic order on every rank, writes merged dense weights, and omits only the adapter keys belonging to those modules. Live parameters, identities and trainability are preserved. CPU export snapshots now own their storage, and ambiguous tied LoRA bases or normalized-key collisions fail explicitly rather than silently dropping an update.

Validation covers each requested point:

• 23 CPU tests passed, including real local HF save/reload with nonzero LoRA and RSLoRA adapters, BF16 bases/FP32 adapters, dtype tolerances, wrappers, repeated exports, and ordinary dense exports.
• Both distributed cases passed on two H100 GPUs using actual FSDP DTensors and NCCL collectives. They check identical collective order across ranks, exactly one gather per parameter, unchanged live weights, and dense HF reload output agreement for both LoRA and non-LoRA models.
• A dedicated two-rank GPU CI step runs those distributed cases; Ruff, both pre-commit configurations, Pyrefly and workflow YAML checks passed.

The tests use a tiny randomly initialized Qwen2 model and local safetensors; this is not a full WTS/VLM training rerun or remote-upload qualification. The inherited training fix is in its own follow-up, keeping #291 focused on source installation.

[lfengad]

Approving this PR. Please track the automatic LoRA HF-export merging gap raised in #291 (comment) as a separate follow-up; it is inherited from the base and is not a blocker for this source-install/setup change.

[ramanathan831]

Following up on the approval review: #308 is the separate follow-up for the automatic LoRA HF-export gap, targeting main. Its current head is 7a54b49; it is implemented and open for review, not yet merged. #291 remains scoped to source installation/setup.

The follow-up covers the requested behavior: merge LoRA/RSLoRA updates into independent dense export tensors, omit the corresponding adapter-only keys, preserve live training weights and ordinary non-LoRA exports, and keep FSDP collectives ordered consistently across ranks. See the implementation and coverage reply.

I rechecked the saved results against that exact head: the 2026-10-23 CI reproduction passed all 23 focused CPU cases and both two-rank H100 export/reload cases. The broader GPU-workflow selection recorded 769 passes; its sole failure was the separate vision_sft_nano CUDA OOM, reproduced identically on unchanged base 9214f43. Existing skips are recorded, and GitHub-hosted Actions still await approval. No additional code change to #291 is needed to satisfy this follow-up request.