Install torch before diffusion learner dependencies - #8
Open
AlexEisie wants to merge 24 commits into
Open
Conversation
…ality check scripts/smoke_models.py launches a short auto-planned run (--budget, no --gpu; auto micro-batch/lora-targets) for every supported model alias, tiered by bf16 footprint so a sweep can be sized to a budget. Sequential and self-cleaning: each run polls the runs registry to a terminal state, tears its clusters down, and lands in a results.jsonl + report.md. scripts/compare_diloco.py quantifies the core product claim — async fragment merging does not hurt much vs synchronous training — at a fixed token budget, scored by held-out eval loss. One synchronous baseline learner (--syncer none; --shard fsdp on GPUs = the FSDP2 baseline) vs M learners through the real Rust syncer over localhost, swept across presets (M=2/4, alpha 0, q4 wire, serial rounds, no heloco, strided). DiLoCo arms are scored on the syncer's merged params via yeto-export — the artifact a real run ships — never a single learner's local weights. Both support --dry-run; neither has been executed beyond it (by request). tests/test_smoke_scripts.py covers the pure logic: tiering, selection, cluster-safe names, budget splits, preset coverage, command construction. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
# Conflicts: # syncer/src/server.rs
Add NAVA fine-tuning backend
--learner-gpus g runs each learner under torchrun on its own CUDA_VISIBLE_DEVICES block (learner i owns [i*g,(i+1)*g)); --shard fsdp shards a frozen base too big for one GPU. Token accounting scales by the ranks per learner. Eval loads bf16 on accelerators (a 10B+ fp32 base would not fit one GPU). Enables e.g. gemma4 arms on a 4-GPU G instance: baseline on 2 GPUs vs M=2 learners x 2 GPUs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The in-process eval kept the base model resident on GPU 0 (25 GB for gemma4), starving the next arm's learner on the same GPU block — found on the first 4xL40S run. Each eval now runs as a child process (--eval-only) whose exit returns the memory. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
subprocess.run returning does not mean the driver has released the dead eval child's VRAM; the next arm's learner raced that window and OOM'd (25 GB still resident on GPU 0). Every eval and arm now blocks until no compute process holds >2 GB, and fails loudly with the offending pids after 180s. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The drain guard parsed per-process memory numerically, so drivers that report [N/A] slipped a fully-resident process (the eval child's 25 GB) past it and the next arm OOM'd anyway. Any listed compute app now counts as a holder, the guard logs what it is waiting on, and evals drain again after the child exits. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
--min-round-interval-ms puts a floor on time between round launches. WAN latency spaces merges naturally (H = many inner steps per fragment); on LAN/localhost rounds complete as fast as learners answer, collapsing H to ~2 and over-driving the outer optimizer (7x effective gain from Nesterov momentum on correlated deltas — measured +33% eval loss vs sync on gemma4/ Lean at H~2). The throttle lets local runs emulate the design-point cadence; paced-launch covered by an integration test (lower-bound assert). compare_diloco gains the arms that separate the hypotheses: 'avg' (merge reduced to plain weighted parameter averaging — frequent averaging should track the sync baseline if the gap is operator gain), 'm2h24' (stock DiLoCo throttled to H~24), and --baseline-loss to reuse an already- measured baseline instead of re-paying its 66-minute run. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
NVSwitch systems need nvidia-fabricmanager running before CUDA will initialize (cudaGetDeviceCount -> Error 802); sky's pinned AMI does not ship it. Found live on a p4d.24xlarge: every CUDA process failed at init. The DL Base GPU AMI (already resolved via SSM for B200) ships the fabric manager preinstalled and enabled, so A100/H100/H200 now route to the same resolver. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A100 now routes to the DL Base AMI resolver (fabric manager); the 'provider default is fine' case moves to A10G. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The m2h24 preset's 56s floor assumed L40S step times; on A100+NVLink the same run finishes before each fragment merges twice (measured: one merge per fragment, eval worse than the untrained base — a single overshooting outer step, not H~24 DiLoCo). The right floor is H*step_time/P, which is hardware-dependent, so make it settable. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
--sync-interval-steps (default 24, the paper's design point) adapts the round-launch floor to the measured learner step time (H*xi_step/P), so each fragment re-merges after ~H inner steps of the slowest learner. On a WAN the natural round latency already exceeds the floor and nothing changes; on LAN/localhost this prevents the measured H~2 pathology (+9% eval loss vs synchronous from outer-optimizer overshoot on correlated deltas; H~24 measured to match synchronous within +0.5% on gemma4/Lean). The comparison harness pins H explicitly per arm (off for the H~2 probe arms, 24 for m2h24) so sync frequency stays an experimental variable. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The NAVA learner's WarmupCosineAnnealingLR was stepped once per local optimizer step, so its phase was a function of this process's inner-step count since launch. On a Spot preemption the relaunched learner reset that count to zero: it re-ran warmup and the full cosine from the top, and the two learners drifted onto independent schedules. Re-base the schedule onto the syncer's global_step — the outer-step clock that is broadcast and re-acquired on every (re)connect, is durable on the on-demand head, and is identical across learners. A relaunched learner now resumes its LR at the correct phase; both learners always share one schedule. - learner: add --total-steps (outer-step horizon T); parameterize the cosine over global_step in [0, T]; replace per-step sched.step() with set_lr_from_clock() at each sync boundary (and a seed before the loop). --syncer none falls back to the local-step clock. Log the current lr. - backend: pass --total-steps through to the learner command. warmup_steps and the cosine horizon are now in global/outer-step units. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
NAVA learner: drive the LR schedule off the syncer global_step
The batch auto-probe (resolve_micro_batch_size) was wired only into the LM learner; the NAVA learner read batch_size straight from config, so "auto" sizing was unavailable for NAVA fine-tunes. Add yeto/nava/autobatch_nava.py and --nava-batch-size auto. Because NAVA's AudioVideoDataset batches internally, the probe rebuilds a one-batch loader per candidate and runs a real fwd/bwd/step, doubling until OOM. It is deliberately conservative so a batch that just fits at probe time does not OOM later on unseen peaks (allocator fragmentation, variable clip resolutions, the sync boundary, activation-checkpoint recompute): - probes under a VRAM cap (YETO_NAVA_PROBE_MEM_FRACTION, default 0.80 -> ~20% slack), then lifts the cap so training runs with the full card; - runs 2 steady-state iterations per candidate; - bounded by a low ceiling (YETO_NAVA_MAX_MICRO_BATCH, default 8). Grad sync is disabled during the probe (ranks pull different real batches, so a one-rank OOM must not hang the island on DDP's backward all-reduce) and consensus is a MIN all-reduce; grad-accum is rebalanced to hold the effective batch. Both tuning vars are forwarded submitter->learner in backend.py. Also add --nava-predownload-uris + yeto/nava/pre_download_clips.py: fetch a list of clip URIs into the resolver cache (stable_cache_path) during learner setup, before torchrun, so the first training step reads from local disk instead of streaming from S3. Opt-in; default keeps same-region lazy caching. backend.py stays torch-free (local _int_or_auto, matching backends/lm.py) so the CLI/launcher path still imports without torch. Adds tests/test_nava_autobatch.py (doubling, ceiling, headroom reserve/restore, OOM-at-1 fallback, non-OOM abort, explicit/CPU short-circuits). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The test call sites passed 8 positional args, omitting the `params` slot before `opt`, so every call raised TypeError (missing `global_step`). Add the missing `params=None` argument to all call sites. resolve_nava_micro_batch's signature is (args, build_loader, pipe, params, opt, cfg, device, world, global_step). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace NAVA special-case backend with generic diffusion components
Document diffusion component model aliases
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
2 participants
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.
Summary
TORCH_SETUPstep for diffusion learners beforerequirements.txtContext
This was found while starting the G4 NAVA LoRA smoke test: diffusion setup needs the same torch pre-install path as LM because
requirements.txtintentionally does not install torch.Validation
PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. pytest -q -p no:cacheprovider tests/test_backend_components.py tests/test_torch_setup.pygit diff --check