More bug fixes and performance improvements - #92
Open
ndryden wants to merge 53 commits into
Open
Conversation
The benchmark driver no longer expands list-valued config parameters into a
cross product of runs: each `scaffold benchmark` invocation now runs exactly
one worker in the resolved run directory, so the per-combination `param_set_i`
subdirectories are gone and restart/checkpoint paths stay in one place. Config
validation now always rejects list values for scalar keys, naming each
offending key and its value ("problem_scale: parameter sweeps are no longer
supported; got list [6, 7]") instead of the previous TypeErrors from
`math.floor(list)` / `int - list` deep inside Config; load_config's config type
"sweep" is renamed "benchmark". This supersedes findings R11, R12 and R16 of
the round-2 review, which all stem from the half-finished sweep path.
A resume whose checkpoint already covers config.epochs leaves the epoch loop before any epoch body runs, so val_loss_avg was never bound and the final-save block crashed with UnboundLocalError on every rank. train() now logs that there was nothing to resume and returns cleanly, leaving the existing checkpoint (which already covers those epochs) alone. Round-2 review: R01.
…rged The training loop runs while the validation dice is below target_dice but that score was never checkpointed, so every restart of a converged epochs:-1 run re-trained one full epoch before rediscovering it had converged, inflating the epoch count and the FOM's total train time. The score now rides along in the checkpoint extras and seeds the loop variable on resume. Round-2 review: R06.
cleanup(train_from_scratch=True) deleted the checkpoint files but kept best_val_loss and last_saved_epoch, so the fresh run's is_best decisions stayed gated by the deleted run's best and it wrote no best checkpoint until it beat a score nothing backed any more. Round-2 review: R02.
…model The existing NaN check tests the hard-argmax dice, which stays finite even for an all-NaN model, so it could never fire; a diverged run kept looping below target while NaN weights overwrote checkpoint_last.pth. The reduced train and validation losses are now checked right after the data-parallel reductions -- identical on every rank, so all ranks raise together. Round-2 review: R03.
wait_for_save now re-raises the writer's exception instead of logging and dropping it, train() consumes the run's final save so a failed async write can no longer exit 0 with no checkpoint, and save_checkpoint broadcasts rank 0's outcome (success or error sentinel, including a deferred async failure) before anyone raises -- previously rank 0 raised ahead of its own broadcast and left the peers in an unmatched collective reporting a transport error instead of the disk error. Round-2 review: R04, R05.
A killed write strands a full-checkpoint-sized checkpoint_*.pth.tmp.<pid> and a quarantined checkpoint keeps its .corrupt copy forever; nothing ever removed either. The from-scratch cleanup now deletes both, and manager construction sweeps orphaned temp files (skipping this pid's) since the kill/restart cycle that creates them always takes the resume path. Round-2 review: R07.
The module had no import sites anywhere in the package, tests or scripts, its usage comments reference a train.py flow that no longer exists, and compare_state_dicts2 crashes on any real optimizer state dict (it truth-tests an elementwise tensor comparison). Round-2 review: R10.
The nothing-to-resume warning suggested lowering target_dice, which is exactly backwards for the converged case, and read oddly for a fresh run that never entered the loop. Report the actual inputs (start epoch, epochs, starting val dice vs target) instead. Round-2 review: R01, R06 follow-up.
Any failure in the rank-0 reuse/generate decision or the final meta-write and rename is now broadcast as an error sentinel, so peers raise the same error instead of hanging in bcast/Barrier. The reuse scan also skips .tmp_* staging dirs and tolerates unreadable metadata, and staging dir names carry pid+uuid so same-second jobs cannot collide. R27
meta.yaml is now written to a temp file, fsynced, and renamed into place, so a killed job cannot leave a truncated document. The loader treats only a missing meta.yaml as legacy v1; a present-but-unreadable or version-less one raises instead of silently reinterpreting a modern dataset. R28
Killed generations left their .tmp_* staging trees under the config_id base forever, so every retry stacked another copy. Rank 0 now removes staging dirs that have sat untouched past an age threshold, which keeps a concurrent job's (far younger) staging dir safe. R37
Categories and instances are derived from config.seed but resume is a pure file-existence check, so a run under a new seed silently adopted another seed's library and published a dataset stamped with the wrong seed. Library paths now carry seed<seed>, so cross-seed reuse is impossible; old-layout libraries are not found and are regenerated. R29
category_search derived its loop-gating remaining count from a per-rank filesystem scan, so ranks with divergent views could post mismatched collectives (bcast against reduce) and hang. Rank 0 now scans and broadcasts the existing-index list, mirroring the instance.py work-list fix. R30
A category CSV truncated by a killed job kept its six-digit name, so the resume scan counted it as done forever while instance generation and the search's own resume both died parsing it. Categories are now staged under a temp name that no scan matches, fsynced, and renamed into place. R31
The centering offset subtracted an extra voxel, biasing every cloud toward the origin: the first half-voxel of each filled axis floored to -1 and was clipped onto plane 0 (1.5x the interior density, with the far plane at 0.5x). DATASET_FORMAT_VERSION is bumped so misregistered datasets are regenerated rather than reused. R33
The guard compared the class count against the int16 limit, but v2 masks ship raw category ids and a sparse split lists only the categories it contains, so a two-entry table holding id 40000 passed and then wrapped negative. The bound is now the largest id the carrier will hold, still the remapped count for legacy datasets. R34
_git_commit_short ran git in the process working directory, so meta.yaml, the published directory name, and the commit-based reuse gate carried whatever repo the job was launched from. It now runs git in the package directory; a non-checkout install still degrades to no-commit-id. R35
Nothing validates vol_size against dc_num_shards, so 16 over 3 shards is accepted as 6/6/4 and per-shard pooling silently diverges from the global result. The fix belongs in DistConv, so this records the hazard as a strict xfail rather than working around it. R32
cli.py decides run dirs and restart state on MPI rank 0 while training uses the launcher's rank environment; a plain torchrun makes mpi4py a singleton in every process, so each claims its own run dir and the job hangs. Cross-check the two world sizes at the CLI entry and abort with an explanation of the launcher wiring. R13
A run launched with --config=PATH emitted a literal --config=__CFG__ into restart.sh, because placeholder substitution only matched whole tokens; the restart then died with "Config file '__CFG__' not found". Substitute the value half of any --flag=PLACEHOLDER token as well. R14
The CLI holds MPI.COMM_WORLD but did not pass its size, and the generator's environment sniffing missed launcher variables the rank side honors (MV2, PALS, and bare SLURM/FLUX task counts), so e.g. a Cray PALS job got a single-process restart.sh. Pass the communicator size and share one world-size variable list with get_world_size. R17
Every rank stat-ed the shared filesystem for a checkpoint and raised on its own verdict, so a divergent view (stale attribute cache) either strands the peers in benchmark.py's timeout-less barrier or lets them run on after rank 0 aborted. Make it a rank-0 decision broadcast to all ranks, matching the other CLI decisions. R18
generate_fractals ran the CLI's rank-0 block too, littering base_run_dir with a timestamped benchmark directory whose restart.sh replayed generate_fractals with --restart/--run-dir, flags that subparser rejects (exit 2). Create the run dir, config dumps and restart script only for benchmark. R19
datagen_batch_size and verbose are accepted config keys, but Config never stores them, so the merge silently replaced them with argparse defaults. Carry the file's values into the merged config and let only options actually given on the command line override them: CLI flag > config file > argparse default. R20
mem_stats called torch.cuda.current_device() unguarded, so gather_and_print_mem -- which BaseTrainer.__init__ invokes unconditionally -- killed any CPU/gloo run launched with -v. Report the missing device and log a fallback instead; the early return is uniform across ranks, so no collective is skipped on one rank only. R21
Copying the base config into the run dir kept its original filename, so a base config named config.yaml overwrote the merged config.yaml the CLI had just written -- and restart.sh points -c at that file. Always copy it to base_config.yaml. R22
Config accepted any unet_bottleneck_dim, so a value outside 0..problem_scale-1 built a U-Net with zero or too many pooling levels and failed much later with an opaque max_pool3d size error naming no config key. Reject it at config time, in Config and again after the CLI applies overrides. R25
export_chrome_trace runs before the barrier that precedes rank-0 post-processing and raises when the run had zero profiled steps (or the write fails), killing the profiling rank and blocking every other rank in that barrier until timeout. Move it into a helper that logs the failure and returns. R15
get_local_size ignored LOCAL_WORLD_SIZE/PMI_LOCAL_SIZE/PALS_LOCAL_SIZE that get_local_rank honors, so it returned 1 and the per-node profiler gate selected every rank while the trace name claimed one node per rank. Add the missing variables, round the node count up, and write the trace into config.run_dir instead of the working directory. R23
The gate enabled profiling for any value except the literal "off", so PROFILE_TORCH=0, =false, =no and ="" all turned the profiler ON -- the opposite of what they say, and inconsistent with the sub-option flags in the same module. BEHAVIOR CHANGE: profiling is now enabled only by 1/true/on/yes (any case); every other value, including previously-enabling ones such as "enabled", leaves it off. R24
The profiler context wraps checkpoint cleanup and every warmup batch while prof.step() advances only per training batch, so PROFILE_TORCH_WAIT=0 buffered that entire prologue in host memory as one unbounded step. Clamp wait to 1 and say so, rather than reordering the context or stepping it from warmup. R26
Regression guard for the merge order: a restart driven by the run dir's config.yaml (what restart.sh emits) must reproduce the first run's CLI overrides and auxiliary config keys and still take the resume path. R20 R22
Rank 0 is the only rank that touches the run directory, and it does so while its peers are already committed to the next collective, so every one of those windows has to report failure *through* that collective. Three were left open. cleanup() ran _remove_checkpoint_files between the drain and the broadcast: individual unlinks were tolerated, but the glob and stat around them can raise on a shared filesystem (ESTALE, EACCES), re-creating the R05 hazard in a narrow window. load_from_checkpoint() folded only the drain error into its decision broadcast, leaving _select_and_load -- which stats, deserializes and renames -- able to kill rank 0 before the broadcast the peers were waiting in. Both now travel as decisions; cleanup's payload carries the phase so the peers name the operation that actually failed rather than reporting a "save". __init__'s orphan sweep is rank-0-only outside any collective, so a directory it cannot list would abort rank 0 alone and strand the peers at cleanup's first collective. It only reclaims space, so it is now best-effort with a warning. VA-1, VA-2, VA-3
The category search broadcasts its scan of the existing category files, but rank 0 then loaded those files' parameters -- for its duplicate guard -- just *after* that broadcast, unfenced. A CSV that will not parse (ragged, or hand-edited) killed rank 0 while every peer had already taken the broadcast and entered the work loop, stranding them in its next collective: the same window class the broadcast was introduced to close. Scan and load are now one guarded decision reported through one broadcast, as get_dataset reports its selection. The datagen consensus guards caught (Exception, SystemExit), which is not the same as "everything": a KeyboardInterrupt delivered to rank 0 alone (Ctrl-C on the launching terminal, a watchdog SIGINT) unwound straight past the broadcast and hung the peers. They now catch BaseException. SystemExit keeps being converted -- it must never escape get_dataset, since a peer would read the silent unwind as success -- while the rank that was interrupted re-raises the interrupt after posting the sentinel, so it keeps the operator's exit status and its peers still learn to stop. volumegen's guard, the same pattern feeding the same allreduce, is fixed with it. VB-2, VB-3
The orphan cleanup judged a staging directory by the newest mtime among the directory and its immediate children. Volume writing lands at depth >= 2 (volumes/<split>/N.npy), and generation is not bounded by the staleness threshold -- at the larger scales it runs for days -- so a perfectly healthy generation stopped touching anything the probe could see and read as "untouched for 48 hours". A concurrent same-config start then rmtree'd it out from under its peers, which died on FileNotFoundError. Liveness is now reported rather than inferred: every rank writing volumes refreshes <staging>/.heartbeat every five minutes, and cleanup keeps any directory whose heartbeat is younger than the threshold. The marker is removed before the staging dir is published, so a dataset does not carry it. Backing the heartbeat up (for staging dirs written before it existed, or killed before the first beat) the mtime probe now walks a bounded number of directory levels instead of one. It stats directories, whose mtimes change when entries are created in them, so it notices a writer without a stat storm over the volume files, and it stops at the first recent entry -- the live case, the one that must not be misjudged, is the cheap one. A tree too wide to walk within the bound is called live: failing to reclaim disk is recoverable, deleting a running job's dataset is not. VB-1
Both atomic writers in the category search name their temp file after the writing pid and unlink it when the write raises -- but a SIGKILL (walltime, OOM, node failure) skips that Python-level cleanup. Nothing ever looked at the strays again, so .NNNNNN.csv.tmp<pid> and .rng_attempt_rank*.tmp<pid> piled up one per killed process. Rank 0 now sweeps them at startup, best-effort, the way instance.py already sweeps its equivalents. The seed-keyed relayout is silent about libraries in the old location: the existence check simply does not find them. From outside that looks like a library that was there yesterday being regenerated for no reason -- hours of work at scale. One warning naming the old directory and the new one explains it. VB-4, VB-6
restart and run_dir can come from the config file as well as the command line -- the generated restart.sh replays a dumped config.yaml, which carries both -- and an absent --restart cannot outrank a file that sets it, since an unset store_true flag is indistinguishable from its default. resolve_run_dir read only the command line while the pre-check read the merged config, so the two disagreed in both directions. A config-file `restart: true` with no run directory therefore created a fresh timestamped directory, wrote its config dumps and restart script into it, and only then died for want of a run dir; it now fails before anything is claimed. And a run_dir inherited from a reused config.yaml stayed in the config while a fresh directory was created to train in, so the pre-check passed on another run's checkpoints and the job died hours later with its dataset generated. resolve_run_dir now reads the merged values, and writes its answer back to benchmark_run_dir, restart and run_dir (cleared for a fresh run), so nothing downstream can re-derive a different one; missing_checkpoint_error keys off the resolved benchmark run dir. The pre-check is also gated on the benchmark subcommand, which is the only one that has run directories at all. VC-1
Rank 0 builds the whole job's config while every peer waits in a barrier, so anything raising in there stranded them: the recomputed bottleneck check and the n_categories check added with R25, the Config() validation that has always been there, and now the run-dir resolution. The user saw a hang where an error message belonged. The rank-0 block moves into build_run_config and runs inside a guard whose outcome is broadcast, exactly like the restart pre-check below it. Rank 0 re-raises the original exception, keeping its traceback; the peers rebuild it from the type name and message that crossed the wire, so a builtin type comes back as itself and anything else degrades to a RuntimeError naming the original. The exception object itself is deliberately not pickled across: a failure to unpickle on the receiving side would turn the error being reported back into the hang it was reported to avoid. VC-2
The launcher-environment helpers called int() on whatever a variable held. A site wrapper that exports WORLD_SIZE= (an unset shell variable) or a placeholder like "auto" therefore killed every scaffold invocation with a bare ValueError naming neither the variable nor a remedy -- and it fired from the top of the entry point, so even `scaffold --help` died. Each lookup now goes through a helper that treats an unusable value as absent and consults the next source in the priority order, warning when the value looked deliberate (as _sniff_launch_shape already did with `if val:`). The Slurm and Flux tasks-per-node divisions no longer trust the node count either: zero is as unusable as a word. The world-size cross-check also moved after parse_args. Asking what the flags are is not a job launch, and answering it with a launcher mismatch -- exactly the environment someone debugging one is sitting in -- helps nobody. It still runs before any run directory is created, which is the property R13 needs. VC-3, VC-4
Flux and Slurm state their node count, so the restart script reproduced their shape correctly. Everything else -- Cray PALS, plain torchrun -- fell back to NODES=1, which relaunched an 8-rank job spread over 2 nodes as NODES=1 TASKS_PER_NODE=8: an oversubscribed node, or a job the scheduler rejects. get_local_size already reads those launchers' per-node variables (PALS_LOCAL_ SIZE, LOCAL_WORLD_SIZE, ...), so the node count is derived from it, ceil-ing the division as the rank side does. required=True keeps "one rank per node" distinct from "nothing reported a per-node count"; only the latter falls back to the historical single-node assumption. VC-5
The "no new epoch was trained" warning covers two ways of entering train() with nothing to do, and described one of them as the other: a fresh epochs:0 run was told "there was nothing to resume", sending its user to look for a checkpoint that was never part of the story. It now says the run had no epoch left to run, which is true of a completed resume and a zero-epoch fresh start alike. (The deeper problem in that scenario -- worker.py's rank-0-only genfromtxt raising after destroy_process_group, so rank 0 exits non-zero while its peers exit 0 -- is pre-existing structure, unchanged here.) explicit_cli_keys' docstring claimed that a flag passed with exactly its default value costs nothing because "both spellings then agree on the default". They do not: --datagen-batch-size 10000 next to datagen_batch_size: 500 in the config yields 500. The rationale is corrected to name the real trade-off rather than deny it; the defaults stay where they are, which is what the R20 tests pin. perf_measure now has a logger instead of printing. Its messages are all "your profiling request was not honored as written", which belongs on a channel a caller can filter or capture, not in the middle of the run's stdout. The R26 test follows it to caplog. And the R21 tests' function-body `import pytest` moves to module level as a skipif marker. The config-error re-raise added a commit ago is spelled as a statement rather than a conditional expression. VA-4, VC-6, VC-cosmetics
test_non_repo_install_reports_no_commit_id builds a "not a checkout" directory under tmp_path, but git walks upwards until it finds a repository: run with --basetemp inside a ScaFFold checkout, that directory inherits the checkout's HEAD and the test fails on where it was run rather than on what it tests. GIT_CEILING_DIRECTORIES stops the walk at tmp_path. VB-7
An improving epoch pickled and fsynced the identical state dict twice, once for checkpoint_last.pth and again for checkpoint_best.pth. The best file is now copied (tmp + fsync + os.replace) from the last file the same writer just committed, halving the serialization CPU and checkpoint bytes. R08
snapshot_training_state cloned model and optimizer state device-to-device, holding ~3x parameter bytes of accelerator memory across the whole warmup phase (and warmup's own memory peak). It now copies to CPU; load_state_dict puts the state back on each parameter's device on restore. R09
instance.main re-ran np.genfromtxt on the same category CSV for every (category, instance) work item -- 145 identical shared-filesystem reads per category. Work items for a category are contiguous in the block partition, so a one-entry cache gives one parse per category per rank. R36
Warmup ran only leading, full-size batches, so with local_batch_size>1 and an indivisible shard the epoch's partial batch met cuDNN/MIOpen for the first time inside the first timed epoch -- a measured 95 s stall in epoch_duration, the FOM denominator. Warmup now runs one extra iteration per distinct ragged size (agreed across ranks, since validation shards are unpadded) inside the existing snapshot/restore envelope. R41
pyplot retained every figure the run-summary plots created, and a sweep calls standard_viz.main in-process once per combination, so figures (and their canvases) piled up for the whole sweep. Each is now closed after its savefig, and unconditionally in a finally since plotting errors are swallowed. R43
The U-Net's use_checkpointing (made correct by F46) had no config key and no production caller, so the memory/compute trade it offers could not be taken. Per the user's decision to wire the feature rather than delete it, an activation_checkpointing 0/1 flag now enables it in worker.py, before the DDP wrap hides the method behind .module. R44
ATen's GroupNorm computes its per-group statistics with a kernel that launches one workgroup per (batch, group) row. At the benchmark defaults (local_batch_size=1, group_norm_groups=8) that is 8 of an MI300A's 228 CUs, and GroupNorm was 86.98 ms of the 186.97 ms scale-7 training step -- 46.5% of wall time, the single largest cost in the step (R38). FastGroupNorm subclasses nn.GroupNorm and routes its forward through a lazily built, process-wide torch.compile(F.group_norm, dynamic=False, fullgraph=True), which hands the reduction to Inductor to tile across the whole device. The module holds the same weight/bias under the same names with no new buffers, so checkpoints round-trip in both directions against a plain nn.GroupNorm build. The compiled path is used only where it is safe, and every rejection falls back to the stock kernel: CPU tensors (the CPU suite pays no compile latency), tensor subclasses such as DistConv's DCTensor (Dynamo cannot trace __torch_dispatch__ wrappers), an already-compiled enclosing region, an explicit SCAFFOLD_GROUPNORM_COMPILE=0, and -- permanently, with one warning -- any exception out of torch.compile. Dynamo's per-function recompile limit is raised from its stock 8: one UNet needs 10 cache entries (5 distinct GroupNorm shapes, each again under no_grad for evaluation), and past the limit Dynamo gives up and silently reverts every GroupNorm to the slow kernel. Measured on one MI300A with review/round2/repros/perf/step_bench.py --layout cl (1x3x128^3, layers=4, bf16 autocast, GradScaler disabled, warm MIOpen db), median of 20 steps, eager numbers from the same build with SCAFFOLD_GROUPNORM_COMPILE=0: step 184.69 ms -> 100.71 ms (1.83x) forward 104.04 ms -> 30.64 ms backward 75.59 ms -> 64.66 ms peak alloc 9.80 GiB -> 8.22 GiB --batch 2 279.58 ms -> 187.61 ms (1.49x; no regression at B>1) torch.profiler over the same step: GroupNorm 86.98 ms/step (46.5% of wall) -> 6.56 ms/step (6.7% of a 97.91 ms step). Compilation is one-time: the first compiled step takes 15.15 s with a cold Inductor cache and produces 5 graphs, after which 12 steps at 97.4 ms produce none; the first no_grad forward adds the other 5. The default 64 warmup batches absorb it outside every timed epoch. Determinism: no gate needed. Two separate processes running three fwd+bwd+Adam steps of the scale-7 UNet under the more_determinism settings (use_deterministic_algorithms(True, warn_only=True), cudnn.benchmark=False, fixed seeds) hash bitwise identically with the compiled path, exactly as they do with the eager one, so no config flag is plumbed through. Not yet visible in the default configuration: worker.py wraps activations in DCTensor even at dc_num_shards=[1,1,1], and GroupNorm then keeps the eager kernel. Verified through a worker.py-shaped DistConvDDP harness that all three paths (DCTensor, plain tensors, plain + use_checkpointing) train without error and agree on their losses; the speedup lands as soon as the unsharded wrap is skipped (R39).
A rank whose warmup fetched no batch returned before the ragged-size all_gather, skipping a collective its peers post — the divergence class this round closes elsewhere. Latent today (padded training shards make loader lengths rank-invariant), but the collective pattern must not depend on local loader state. Found by the final verification pass.
The R08 comment claimed the copy halves the bytes pushed at the shared filesystem; measured on Lustre the copy is ~1.02x — the real saving is the serialization CPU. Note the redundancy trade-off (best is now a byte copy of last), use a 16 MiB copy buffer to cut syscall count on parallel filesystems, and document that the Dynamo recompile-limit raise clobbers a deliberately smaller limit. Found by the final verification pass.
The seed-keyed relayout (R29) made the shipped seed-unknown CSVs at ScaFFold/fractals/var0.15/3DIFS_param unreachable under any configuration. Remove them, their package-data glob, the dead Config.library_root (zero readers), and the README claim; libraries regenerate deterministically from the configured seed. Resolves verification item VB-5 per user decision.
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.
Another round of fixes from Claude Fable 5.
Note that I expect very modest, at best, performance improvements; the headline improvement (groupnorm) isn't really applicable yet and may get superseded.