Many bug/performance fixes#88
Open
ndryden wants to merge 27 commits into
Open
Conversation
… helpers - pyproject.toml: add psutil to runtime deps (imported unconditionally by worker.py), a 'test' extra, and pytest configuration (testpaths, mpi/gpu/slow markers, global 120s timeout). - tests/conftest.py: shared fixtures for a minimal RunConfig, v1/v2 dataset directories built without running datagen, run-directory skeletons, a CPU PyTorchTrainer, and fresh-subprocess execution (needed for numba RNG tests). - tests/helpers/mpi_runner.py: subprocess harnesses that convert MPI/torchrun deadlocks into per-rank stack traces via faulthandler; mpi_run skips when no launcher is available, torchrun_gloo runs multi-process gloo tests locally. - tests/test_infra.py: smoke tests for all fixtures and harnesses.
- standard_viz: create the figures directory idempotently (re-runs and resumed runs previously crashed rank 0 with FileExistsError), wrap plotting in try/except so a plotting failure can never kill an otherwise-successful run, save the validation-Dice curve as val_dice.png (it was previously written to val_loss.png), and plot the actual validation loss into val_loss.png. - utils.plot_img_and_mask: label each mask panel with the class index actually rendered (labels were off by one) and coerce the class count to int. - datagen/data_visualizer: render 4D channels-first (3, D, H, W) volumes by computing an occupancy grid and passing per-voxel facecolors to ax.voxels; keep the 3D path for masks; reject other ranks with a clear error. - perf_measure: enable the torch profiler when PROFILE_TORCH is set even if Caliper was requested but failed to import (previously silently ignored). - tests/test_reporting.py: regression tests for all of the above.
- UNet.__init__ now rejects layers < 1 with a clear error. Previously a model with problem_scale == unet_bottleneck_dim built silently and crashed with an IndexError at the first forward pass, with the bottleneck block never invoked. - use_checkpointing() previously called the torch.utils.checkpoint module as a function (TypeError) and assigned non-Modules into ModuleList. It now sets a flag consumed in forward, wrapping each block in torch.utils.checkpoint.checkpoint(..., use_reentrant=False). Off by default; outputs and gradients are identical either way. - tests/test_unet.py: build-time validation, forward-shape regression guard, and checkpointing output/gradient equivalence tests.
… shape - restart.sh previously hardcoded site-specific module loads, an LD_PRELOAD path, and an unconditional 'export PROFILE_TORCH=ON', silently enabling the profiler on restarted runs and aborting on machines without those modules. The script now re-exports PROFILE_TORCH/CALI_CONFIG only when they were set in the generating run, derives venv activation from the active VIRTUAL_ENV (non-fatal if missing), and emits no foreign environment. - The single- vs multi-rank template choice was gated on FLUX/SLURM env sniffing only, so mpirun- or torchrun-launched multi-rank jobs got a single-process restart script. create_restart_script now accepts an explicit world_size and also recognizes WORLD_SIZE, OMPI_COMM_WORLD_SIZE, and PMI_SIZE; the multi-rank template is chosen whenever world_size > 1. - tests/test_restart_script.py: content tests over generated scripts, plus a bash -n syntax check for every variant.
- generate_fractal_points returns an empty (0, 3) array for numpoints <= 0 instead of writing out of bounds inside the njit kernel, and computes the runaway check in-loop (any |coord| > 1e10) instead of hard-coding the flag to True, making the downstream quality checks and counters live again. - tests/datagen/test_ifs_kernel.py: unit tests for zero-point safety (bounds-checked in a subprocess) and runaway detection.
…sistent - BasicDataset built its id list from unsorted os.listdir, so the index-to-file mapping depended on filesystem enumeration order. Under DistConv spatial sharding, ranks in one replica resolve the same sampler index independently; divergent readdir views (e.g. across parallel-filesystem clients) would silently stitch shards of different volumes into one training sample. The id list is now sorted, and when a torch.distributed process group is live, dataset construction all_gathers a digest of the id list and raises on any cross-rank mismatch instead of training on corrupt samples. The check is skipped when no process group is live (e.g. plain dataset construction in tests or tools). - Legacy (v1) datasets remapped mask voxel values through per-split mask_values tables, so train and validation assigned different class indices to the same raw value whenever a category was missing from one split. v1 loading now builds one global remap table from the union of all *unique_mask_vals pickles in the dataset root. v2 dataset behavior is unchanged. - tests/test_data_loading.py plus a 2-rank gloo rank script covering both fixes and a v2 regression guard.
…est selection - Checkpoints were written in place with torch.save directly onto checkpoint_last.pth, so a walltime kill mid-write truncated the new file and destroyed the previous good one; load_from_checkpoint then raised with no fallback, permanently breaking restart. Writes now go to a temp file in the same directory, fsync, then os.replace, for both last and best. On load failure the corrupt file is renamed aside (.corrupt) and loading falls back to checkpoint_best.pth, raising only if nothing loads. - Save exceptions were swallowed with a print-and-continue; they are now logged with a traceback and re-raised. - With async_save, is_best was decided by torch.load-ing checkpoint_best.pth while the background writer might still be replacing it. The best val loss is now cached on the manager (seeded from disk once at construction, refreshed on load) and the per-save disk probe is gone. - _transfer_dict_to_cpu aliased tensors already on CPU, letting the async writer serialize live training state; CPU tensors are now cloned. - The final/converged model was only checkpointed when the last epoch happened to be a multiple of checkpoint_interval; train() now writes a final checkpoint on loop exit whenever checkpointing is enabled and the last completed epoch was not already saved. - GradScaler-skipped steps (inf/NaN grads triggering scale backoff) no longer advance the optimizer-step counter: a scale decrease across update() marks the step as skipped, and global_step now counts applied optimizer steps. - tests/test_checkpointing.py: coverage for all of the above.
- The reported validation Dice was computed on softmax probabilities, so it measured calibration rather than segmentation quality: a model with a perfect argmax but modest confidence scored well below 1.0, and any validation sample missing a foreground class contributed a near-zero per-class score because softmax outputs are strictly positive. Since val_score is the training-termination gate (target_dice), a perfect model could never reach the default 0.95 target on datasets with absent classes. The score is now computed from one-hot argmax predictions (per-voxel, so DistConv sharding is unaffected), which also makes the empty/empty guard in the dice computation reachable. Softmax probabilities still feed the loss term, mirroring the training objective. - evaluate() now rejects n_categories < 1 instead of silently reporting Dice=1/loss=0 through the degenerate single-channel path. - val_loss_avg is now a sample-weighted mean instead of a mean of per-batch means, so a smaller final batch is no longer overweighted. - tests/test_evaluate.py: stub-network tests for score hardness, confidence invariance, error-rate correspondence, the degenerate-config guard, and batching invariance of the loss average.
- generate_fractal_points draws from numba's internal RNG, which Python-level np.random.seed never touches, so instance generation was non-deterministic: the same config and seed produced different data on every run. A new datagen.rng module provides a jitted seed_numba() plus derive_seed(), an unsalted BLAKE2b-based mixer producing 32-bit seeds from integer keys. Instance generation now seeds per work item on (seed, category, instance), so outputs are independent of MPI world size, rank assignment, and resume state, mirroring volumegen's per-volume seeding. - Resumed or extended category searches replayed the identical NumPy parameter stream from np.random.seed(seed + rank), re-accepting and re-saving the same IFS parameters under new class indices — bit-identical duplicate classes, silently corrupting the segmentation task. Candidate parameters now come from SeedSequence((seed, rank, attempt_index)) with the per-rank attempt counter persisted atomically alongside the results, so a resume continues the stream instead of replaying it. A save-time dedup guard additionally skips candidates identical to an existing category CSV. - The next category index was the count of existing files, so a numbering gap caused the run to overwrite an existing category. Indices are now derived from parsed filenames (fill the lowest hole, else max+1) and saving refuses to overwrite an existing file. - tests/datagen/test_rng_determinism.py: fresh-subprocess determinism tests, rank-layout independence, resume-stream continuation, dedup, and gap-filling index allocation.
…ation - Weighted instance generation could turn a category validated as contractive into an expansive map (weight rows scale IFS matrix columns up to 1.6x), and the resulting NaN/inf point clouds were saved unchecked; rasterization then clipped NaN-cast indices into the ground-truth mask as garbage voxels under a legitimate class label. Instances are now validated with the runaway flag plus a finiteness check, retried with attenuated weights, and finally generated unweighted so every weight slot fills deterministically; points_to_voxelgrid rejects non-finite input outright. - Instances were written with np.save directly to their final filename and resume accepted any file matching the name glob, so a job killed mid-write left truncated .npy files that were accepted forever and only crashed at training time. Saves now go through a temp name the resume glob cannot match, with fsync and os.replace; stale temp files are cleaned at startup and resume validates candidates, regenerating unreadable files. - mask_detection picked glob(...)[0] blindly: extensionless ids crashed with IndexError and duplicate stems silently used an arbitrary file. Exactly one match is now required, with an error naming the id and the matches. - Volume rasterization normalized each axis by its own extent, stretching every fractal to fill the grid and destroying aspect ratio. Normalization is now isotropic (max extent) with the occupied region centered. The dead 'scale' knob is rejected with a clear error when set to anything but 1 (previously it tripped an opaque shape assert deep in the volume loop). - tests/datagen/test_artifacts.py: coverage for all of the above.
…eeps - Config silently dropped every unrecognized YAML key, so typos and unsupported options vanished without a trace. Config now validates input against an explicit known-keys set (plus the auxiliary keys the CLI layer injects) and raises naming the offending keys; strict=False downgrades to a warning. async_save is now a real, defaulted option instead of an unreachable getattr. The load_config error message also interpolated the builtin 'type' instead of the config_type argument. - A second --config on the benchmark subcommand previously replaced the first wholesale (argparse last-wins) and could not express partial overrides. --config is now action='append': the first file is the complete base config and later files are partial overrides merged and validated via the new load_config_files; unknown override keys are rejected by name. - The benchmark subcommand advertised one run per parameter combination but ran the raw config once, leaking list values into scalar arithmetic. List values in scalar-typed fields are now expanded via a cross product into one scalarized run each, with per-combination run directories; genuinely list-typed fields (dc_num_shards, dc_shard_dims) are never sweep dimensions, and a list reaching a scalar field in a single-run config fails fast naming the key. - tests/test_config.py: validation, override-merge, round-trip, and sweep tests.
- get_dataset's error consensus could never complete: the gather of error messages ran only on rank 0 (inside the failure branch), so any generation failure left rank 0 blocked in gather and the other ranks in the final barrier until walltime. Every rank now participates in an allgather of error status and all ranks raise with the collected messages; non-root ranks no longer return a dataset path that was never finalized. - The dataset-reuse decision was made independently per rank from a shared filesystem scan, so divergent views (e.g. metadata caching across nodes) could strand ranks in mismatched collectives. Rank 0 now scans and decides alone, broadcasting (reuse, path) or (generate, tmp, dest) before any rank acts. - volumegen killed a single rank with sys.exit(1) on a missing instance file mid-collective-sequence, hanging its peers. It now raises FileNotFoundError; worker exceptions are caught locally, per-rank status is combined with an allreduce (and errors allgathered) so every rank executes the same collective sequence on both success and failure paths. - instance generation derived its global work partition from per-rank filesystem globs; divergent views caused duplicate writes or missed instances. The work list is now built on rank 0 and broadcast before slicing per-rank shares. - tests/datagen/test_mpi_consensus.py: single-process tests drive the real functions with a recording fake communicator to pin the decide-and- broadcast structure and identical ok/error collective sequences; 2-rank MPI tests (marked mpi) skip where no launcher exists and run on a cluster.
- DistConvDDP was built with device_ids=[local_rank] while get_device() intentionally selects index 0 under per-rank GPU visibility masking (the primary documented launch mode), crashing every rank with local_rank > 0 at first forward with an invalid device ordinal. The wrap (now a testable wrap_model_ddp helper) pins device_ids/output_device to the get_device() result -- correct in both masked and unmasked modes -- and passes None on CPU as DDP requires. - Launcher detection: world rank/size now honor torchrun's RANK/WORLD_SIZE (previously only LOCAL_RANK was recognized, so torchrun ranks all reported world rank 0), plus PMI and PALS variables, with a final fallback to mpi4py's COMM_WORLD. - initialize_dist now binds the compute device via get_device() before the first collective (NCCL otherwise guesses rank % num_gpus, breaking non-block placements), passes device_id to init_process_group on CUDA, and selects the backend dynamically (nccl with GPUs, gloo without) so one-rank CPU jobs can initialize. - The debug launch-info log evaluated dist.get_backend() eagerly, which raises when logging happens before the process group exists; it is now guarded. - Rank-0 post-processing (FOM computation, figures) previously ran before the final barrier, so a rank-0 exception stranded all other ranks; the barrier and group teardown now precede it, and a stats file with no epoch rows produces a clear error instead of a NaN score. - tests/test_worker_dist.py: env-detection matrix, device-binding order, dynamic backend, DDP wrap kwargs, and a one-rank gloo end-to-end smoke test of worker.main.
…icas - The CSV train loss and train dice were rank 0's replica-local, unweighted batch means: with N data-parallel replicas they represented only 1/N of the epoch's data and weighted a partial final batch like a full one. The trainer now accumulates sample-weighted [loss_sum, dice_sum, sample_count] on device, reduces across the data-parallel group, and writes global sample-weighted means. - Best-checkpoint selection and the logged validation loss used rank 0's replica-local val_loss_avg, so checkpoint_best.pth could track a replica quirk instead of the true global best. The validation loss now rides the existing dice reduction ([dice_sum, val_loss, numsamples]) over the data-parallel group before it feeds the CSV and save_checkpoint. - Metric reductions use the data-parallel subgroup when DistConv spatial sharding is active: every rank in a replica already holds the identical spatially-reduced value, so reducing over the world group would multiply by shards-per-replica. - The validation DistributedSampler padded ranks by duplicating samples, biasing the convergence-gating val_score (a duplicated easy/hard volume shifts time-to-target). Validation now uses an unpadded sampler with contiguous uneven shards; a rank may legitimately hold zero validation samples, contributing nothing to the SUM reductions, and the zero-batch construction error now applies only to single-replica runs. - Fixed a latent multi-rank deadlock in the final-checkpoint guard: it read the manager's last-saved epoch, which is only recorded on rank 0, so ranks diverged on whether to enter the collective save. The guard now uses a rank-symmetric local epoch tracker. - tests/test_metrics_dist.py: four 2-rank gloo tests (via torchrun) driving a real trainer at tiny scale, covering padding bias, global train metrics, global val loss, and best-checkpoint selection.
…tence - Resume plumbing worked only for the exact pair '--restart --run-dir': '--run-dir' alone crashed rank 0 with AttributeError (benchmark_run_dir was only set on the fresh-dir path) while other ranks hung in a barrier, and '--restart' alone created a brand-new directory and later crashed on run_dir=None. Run-dir resolution is now a testable resolve_run_dir function with fixed semantics: an explicit --run-dir always resumes in that directory (implying restart); --restart without --run-dir is rejected with a clear error on all ranks before any collective; otherwise a fresh timestamped directory is created. benchmark_run_dir is set on every path, and the benchmark restart branch now fills run_dir/run_iter. - Fresh run directories are created with exist_ok=False plus a numeric-suffix retry, so two jobs launched in the same wall-clock second no longer share one directory and clobber each other's config and stats. - Restarting without a checkpoint (the shipped default disables checkpointing while restart.sh always passes --restart) appended a second CSV header mid-file, which parsed as an all-NaN row and made the FOM NaN. The stats file is now truncated before the header is written when training starts from epoch 1, and a real resume rewrites the header if the CSV is missing instead of producing a headerless file. - global_step (which counts applied optimizer steps) now rides the checkpoint extras and is restored on resume, so step accounting no longer resets to zero after every restart. - tests/test_resume.py: run-dir resolution matrix, header hygiene for all three resume shapes, step-counter round-trip, and same-second collision handling.
Trainer construction and evaluation assumed both a DistConv parallel strategy and a CUDA device: DCTensor.from_shard dereferenced a None parallel strategy (crashing the process), evaluate() unconditionally derived a spatial mesh from the strategy, SpatialAllReduce assumed a mesh, minibatch timing used CUDA events on CPU devices, CUDA memory statistics were queried on CPU runs, and the epoch summary formatted a None minibatch time. Tensors now pass through unsharded when no parallel strategy is set, the spatial reductions treat a missing mesh as local-equals-global, and the CUDA-only instrumentation is gated on the device type. A two-epoch one-rank CPU (gloo) run through worker.main (training, validation, checkpointing, FOM, and figures) completes.
instance generation indexed the category parameter directory with a bare sorted os.listdir, so any non-CSV bookkeeping file in that directory (such as the category search's persisted RNG attempt counters) shifted the category-index-to-file mapping and crashed generation. Only *.csv entries now participate, matching how the category search itself scans the directory. Found by the end-to-end pipeline run (generate_fractals -> benchmark -> restart).
Voxelization built a dense vol_size^3 boolean grid per fractal and painted the volume and mask through full-volume boolean-mask assignment. Add points_to_voxel_indices to return the deduplicated integer voxel indices and paint by fancy indexing, so the inner loop does work proportional to the number of occupied voxels rather than the whole grid. points_to_voxelgrid stays as a thin wrapper for any dense-grid caller and is bit-identical to before. Point clouds are now saved and loaded as float32 instead of float64. Coordinates are normalized to roughly [-1, 1] and destined for integer voxel binning, so the extra precision was unused while doubling instance file size and read/write bandwidth. Bump the dataset format version so a float64-era dataset is not silently reused under the new dtype.
The mask-value scan globbed the whole directory once per id (O(N^2) metadata traffic on a shared filesystem). Build the id-to-path mapping a single time from one directory listing and hand each worker its resolved path; ambiguous stems (two files sharing a name) are rejected up front with a message naming the id instead of silently picking one. Category search ran a fixed per-rank batch of attempts every round regardless of how many categories were still needed, so the final round could generate thousands of valid categories only to discard all but a few. Size each round from the categories remaining and the acceptance rate observed so far, capped at the configured batch size and never below one attempt per rank. Rank 0 decides the round size and broadcasts it so every rank stays in lockstep for the gather; each rank advances its persistent attempt counter by exactly the round size, so resumes remain reproducible.
Up.forward always called F.pad on the upsampled activation, and F.pad allocates a fresh tensor (a full copy) even when every pad width is zero. With power-of-two volumes the decoder sizes always match, so the pad was a pure copy of the largest decoder activation on every forward. Guard it behind a nonzero-diff check; output is unchanged whenever any dimension actually needs padding.
Several behavior-preserving efficiency fixes to the hot loops: - One-hot targets are scattered directly into a float32 buffer instead of F.one_hot (which materializes an int64 tensor) followed by permute and a float cast, removing a large int64 intermediate and an extra copy every step. - log_softmax is computed once per step and shared: cross entropy consumes it via NLL and the soft-dice term recovers probabilities with exp(), instead of upcasting the logits twice and recomputing softmax on top of cross entropy's internal log-softmax. - The cross-entropy numerator and its normalizer are reduced together in one spatial all-reduce rather than two. - Memory-stat gathering (which resets peak counters and, at debug verbosity, issues world-wide collectives) now runs only on the first batch of the run instead of every batch inside the timed region. - Gradients are cleared with set_to_none=True, skipping a full-gradient memset each step, and the per-batch dice score is detached so the epoch accumulator no longer retains each batch's autograd graph. - Validation accumulates loss and dice as device tensors and reads them once after the loop instead of syncing the host every batch, and its host-to-device copies use non_blocking=True to exploit the pinned-memory loaders. - The Torch profiler now records a bounded, configurable window (driven by prof.step()) with record_shapes/with_stack off by default, instead of recording the entire multi-epoch run and growing an unbounded trace. - Cross-entropy class-weight estimation loads only the mask (via the dataset's mask-only accessor) rather than reading and preprocessing the full image volume for each sampled index.
Sample lookup globbed both the image and mask directories on every __getitem__, so each fetch cost two full directory scans and an epoch was quadratic in the number of files -- costly on a parallel filesystem. Resolve the id-to-path mapping once at construction (asserting one file per stem) and look it up in constant time per item. Masks are shipped from the loader in a narrow int16 carrier rather than widened to int64 in the worker; the trainer and evaluator already cast to long on the compute device, so this quarters the pinned-memory footprint and host-to-device transfer for masks. In the non-sharded path the volume-preparation copy is now skipped when np.load already returned a fresh, contiguous, correctly-typed array (the materializing copy is kept only for the mmap-backed sharded path). A mask-only accessor lets class-weight estimation read masks without also loading and preprocessing the full image volume.
On resume every rank opened the same checkpoint file and ran torch.load concurrently -- an N-way read of one multi-GB file from the parallel filesystem at job start, exactly when a restart is racing its walltime. Rank 0 now selects and loads the checkpoint once (keeping the last-then-best fallback and the rename-aside of a corrupt file) and broadcasts the deserialized state to the other ranks, which no longer touch the filesystem for it. The single-process path is unchanged, and all ranks return the same resume epoch or fail together rather than leaving peers to hang.
The benchmark now always runs distributed (a one-rank launch is the supported singleton case), so the test suite no longer passes dist=0 configs. An autouse fixture provides the RANK/WORLD_SIZE/LOCAL_RANK environment a one-rank launcher would export, and trainer tests run under a one-rank gloo process group (new gloo_group_1rank fixture) because CE-weight estimation, metric reductions, and checkpoint broadcasts now issue collectives unconditionally. The worker smoke test runs worker.main end to end as a one-rank gloo job (forcing the CPU path so it cannot depend on a working GPU stack) and asserts a real unsharded ParallelStrategy is constructed. The stale nondist variants and the removed benchmark_testing.yml config test are dropped, and the DDP-wrap test asserts integer device indices on CUDA and None on CPU, matching what DDP accepts. Test scaffolding follows suit: the expected train_stats.csv header includes the new optimizer-steps columns, _git_commit_short and _write_to_disk shims take the log parameter, and _write_to_disk tolerates log=None (it is optional for CheckpointManager).
total_optimizer_steps feeds the per-epoch CSV column, the completion message, and the final adiak value, but it was not saved in checkpoint extras or restored on resume, so a restarted run reported step totals that ignored all pre-resume work. Save it alongside global_step in both the interval and final checkpoints and restore it in cleanup_or_resume. The two counters advance in lockstep (both count applied optimizer steps), so a checkpoint from before the dedicated key existed falls back to global_step rather than resetting the total to zero.
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.
Address a variety of issues identified by Fable 5. Approximately one issue per commit, with tests added. Apologies for the large number of changes but one PR was easier.