Optimize GroupNorm - #93
Closed
ndryden wants to merge 62 commits into
Closed
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
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.
worker.py wraps every activation in a DCTensor even at dc_num_shards=1, and FastGroupNorm rejected tensor subclasses outright, so production never took the compiled path: the 1.84x GroupNorm win measured in round 2 was latent. DistConv's __torch_dispatch__ has no GroupNorm-specific handling -- it unwraps to the local shard, runs the stock aten kernel and rewraps -- so there is no distributed GroupNorm semantics to preserve, only a wrapper Dynamo cannot trace. forward() now does that same unwrap itself, in front of the compiled kernel. It cannot use dispatch's mechanism: dispatch runs below autograd where a bare _tensor read is safe, while forward runs above it and must go through DistConv's _ToTensor/_FromTensor pair or the graph back to the producing convolution is severed. Measured on MI300A at the scale-8 GroupNorm shapes, fwd+bwd, DCTensor-wrapped: [1,64,256^3] 129.56 -> 11.82 ms (10.96x), [1,128,128^3] 31.28 -> 3.19 ms, [1,256,64^3] 7.85 -> 1.11 ms; the two smallest shapes are launch-bound and within noise. Values and gradients are bitwise identical to the eager wrapped route, including a 2-rank sharded run. Statistics stay per-shard at every shard count, exactly as DistConv computes them today; making them global is an upstream question (R32/R39), not one this change touches.
Nothing imports this yet; it is inert until a later commit wires it into FastGroupNorm. With PYTORCH_MIOPEN_SUGGEST_NHWC=1, which production sets, every convolution in the UNet emits channels_last_3d but GroupNorm consumes it and emits contiguous -- 22 layout breaks per forward -- and torch's compiled GroupNorm is 6.5x slower on channels-last input than on contiguous. The loss is pure access pattern: in NDHWC one program reads a dense (BLOCK_S, C) run and reshapes the inner axis to (G, C/G), so all groups' statistics come out of one coalesced pass, where Inductor walks the logical NCDHW order over channels-last memory as a strided gather. Measured on MI300A, fp32, 22 GroupNorm sites of a scale-8 UNet, fwd+bwd: 442.8 -> 69.0 ms/step (6.4x), which is a dead heat with compiled GroupNorm on contiguous input while additionally preserving the layout. 95-98% of measured streaming bandwidth at the large shapes. Statistics use a corrected two-pass per tile combined by Chan's parallel formula rather than E[x^2]-E[x]^2, which costs 0.8% at the dominant shape and is 4-23x more accurate than ATen at large input means (at mu/sigma=1e5 the naive form loses the variance entirely). Reductions are register-only with a fixed order and no float atomics, so results are bitwise reproducible. Registered as scaffold_gn::group_norm with a fake kernel and autograd, so it traces under torch.compile(fullgraph=True) without a graph break and composes with DistConv's DCTensor. Only 5-D channels_last_3d input takes the kernel; is_supported() returns False for contiguous input so callers keep their own compiled fallback, which is already at 89-92% of roofline there. Also: fp32/bf16/fp16 with fp32 statistics and torch's autocast contract, an int64 tile-base path for volumes past INT_MAX (measured free, verified at 2.16e9 elements), and an optional fused ReLU that is bit-exact against unfused+F.relu and takes 38% off the forward.
An independent review of 1a8aced attacked the kernel and its suite; these are its findings, plus the edge suite it wrote to pin them (153 tests, none of which the original 51 subsumed). The serious one: every kernel launched on torch.cuda.current_device() rather than the input's device, so a tensor on cuda:1 while cuda:0 was current took a memory access fault and dumped core, where ATen's group_norm handles the same call. One rank per GPU hides this in production, but nothing guaranteed it. The guard costs 0.51 us and is skipped outright when the device is already current. The rest were quieter. The backward fake kernel promised contiguous d_input strides while the real one returns channels-last, which eager never notices and torch.compile miscompiles. mean/rstd were differentiable and answered grad(mean.sum(), x) with zeros -- a plausible wrong number rather than an error -- after materializing a full-size cotangent and running the entire backward for it. (1,C,1,1,1) input was accepted and returned bias where stock raises. Single-element groups returned d_input = 2.2e-05 instead of the exact zero, because fma(dy, w, -c1) contracts and leaves the product's rounding error multiplied by rstd = 316; that case now answers analytically. Double backward and subnormal eps are documented rather than fixed. The review also proposed deleting the Welford correction term, having measured that it does not move the output error. It does not: the fp32 mean rounds the correction away in y. It moves rstd, which is what it computes -- 1472x in a single-tile reduction at mu/sigma=1e6, 3.8x at [1,512,32^3] in the production configuration -- so it stays, now priced honestly at 0.9% of fwd+bwd (the docstring's old 0.8% was the cost of the whole two-pass rewrite) and pinned by a test that measures rstd rather than y. Mutation coverage over both suites: 38/38 killed, up from 30/33. Interleaved A/B over 4 runs shows no fwd+bwd regression at any of the six shapes.
Seven launches per fwd+bwd become four: _stats_finalize is folded into _normalize and _bwd_finalize + _dwdb_reduce into _dx, each elementwise program recomputing the tiny finalize redundantly rather than round-tripping through a launch. Folding alone regresses the dominant shape by 34%, which is the part worth recording: at [1,64,256^3] the elementwise grid is 131072 programs and each would re-read all 2048 partials, so the "tiny" finalize becomes 25.7 GB of reads against a 4.3 GB tensor. Lowering the split count cannot fix it -- the stats kernel needs ~456 splits at that shape just to fill the device. So the elementwise grid is now capped separately (GNConfig.elem_progs) with each program striding over its tiles, which makes the redundant read nprog x nsplit instead of nblk x nsplit: 0.1-403 MB, L2-resident, 1-9% of the tensor. Tiles and split counts were then retuned jointly with the folding, per shape, since the two are not independent. Every candidate was measured interleaved against the incumbent in one process; measured sequentially, this node's background load drifts 35% and swamps the effect entirely. fwd+bwd, min of 7 medians-of-20: [1,64,256^3] -1.7%, [1,128,128^3] -6.8%, [1,256,64^3] -4.9%, and -7.7% to -9.5% at the six small and N>1 shapes. No shape regresses. Scale-8 rollup over 22 sites: 69.5 -> 67.1 ms/step. Results differ bitwise from the previous table -- a different split count is a different reduction order -- but accuracy is unchanged (1.1-1.9e-07 against float64 for both, on all six shapes) and determinism is undiminished: grid, split count, tile sizes and reduction order all remain pure functions of the shape, verified run-to-run and across three interpreters. A new test pins that elem_progs, the one plan field that is a free parameter, cannot move a bit. Also checked and deliberately not changed: the tl.sum(tl.sum(x, 2), 0) reduction in the partial kernels. The documented LDS-staging hazard does not reproduce here (5 shapes x 5 tile sizes x 3 warp counts all compile both ways), and the equivalent spelling measures 0.84% slower, so changing every result's bits buys nothing.
FastGroupNorm now tries Triton, then the compiled kernel, then eager, and a failure at any rung latches that rung off and drops to the next rather than all the way to the bottom. SCAFFOLD_GROUPNORM_TRITON gates it exactly as SCAFFOLD_GROUPNORM_COMPILE gates the rung below. DCTensor is unwrapped explicitly for both fast rungs rather than dispatched through __torch_dispatch__, which keeps the tensor-subclass policy in one place: the kernel's is_supported() only asks isinstance, so an unknown wrapper would otherwise be routed into it where today it stays eager. The ReLU that followed every GroupNorm moves inside it. FastGroupNorm gains an activation argument and always applies the ReLU -- fused into the Triton store, explicit and in-place on the other two rungs -- and DoubleConv keeps an nn.Identity in the vacated Sequential slot so the state dict does not move by one key. That is asserted on the serialized bytes, not just the key list. Scale 7, DCTensor, channels-last, bf16 autocast: 193.25 -> 91.41 ms/step (2.12x), peak allocated 8.33 -> 7.21 GiB. Most of that is not GroupNorm. Its own kernels only go 18.69 -> 7.53 ms; the rest is elementwise and copy time, 99.68 -> 13.63 ms, which was the conv-side layout conversions GroupNorm was forcing by emitting contiguous into a channels-last model. Counted together, GroupNorm and the layout tax it imposed go 118.4 -> 21.2 ms/step, from 61.8% of device time to 23.9%. Whole-model gradient deviation from eager is 4.83e-03 median / 6.73e-02 max over 64 parameters, against 4.94e-03 / 6.64e-02 for the compiled rung that production already ships; triton-vs-compiled is smaller than either. Two processes under more_determinism produce identical losses, activation hashes and parameter hashes, with a counter confirming the Triton rung ran. One fix beyond the wiring: the ladder now re-raises _StopRecomputationError. Non-reentrant checkpointing ends its recompute by raising that from a saved-tensor pack hook, i.e. from inside whichever op is saving a tensor. Absorbing the trailing ReLU moved that boundary inside the try, so the ladder caught it, decided the kernel was broken and dropped the whole model to eager for the rest of the run. The hazard predates this commit; it was simply unreachable while the ReLU sat outside forward().
The one that matters: the fused ReLU silenced NaN. tl.maximum(y, 0.0) maps NaN and -Inf to zero where F.relu propagates NaN, so a diverged run showed a finite forward and a NaN backward -- and round 2 added the non-finite loss abort precisely so divergence stops the run instead of checkpointing a broken model. Testing the complement (tl.where(y <= 0, 0, y)) keeps NaN on the pass-through side. The backward gate had the same defect and matters as much: threshold_backward passes the gradient where the result is NaN, and pre > 0 was zeroing it. Now bit-identical to F.relu on NaN, both infinities and -0.0, on all three rungs, forward and backward. Cost: +0.18% on the 22-site rollup, against a -0.30% noise floor from the activation=None control. A latch flip between a checkpointed forward and its recompute killed the run. Matching the output memory format across rungs does not fix it -- measured: the rungs save different tensor sets, so it still dies comparing a (64,) to a (1,8). The fix is that a latch no longer demotes a module that has already had a call served by that rung, so a block's forward and its recompute always agree; a broken install still costs exactly one attempt per module. Rungs now preserve the input's memory format anyway, since a fallback that returns contiguous re-breaks the channels-last chain this whole line of work exists to protect. _CONTROL_FLOW_EXCEPTIONS is gone. A denylist of framework mechanisms that legitimately raise through a forward was wrong twice (_StopRecomputationError, then CheckpointError) and is unbounded. Narrowing by exception type does not work either -- a HIP launch failure and a CheckpointError are both RuntimeError. So the narrowing is by scope instead: a decorator tags failures raised inside the kernel call itself, which is a closed region that ends before anything is saved for backward, and the ladder catches only that tag. That also removes the retry's saved-tensor-hook asymmetry structurally rather than defensively. The compiled rung catches TorchDynamoException, whose members are all raised at compile time. Remaining: predicates decline under functorch transforms instead of latching both rungs off; a predicate that cannot answer falls back without latching; a transient OOM neither latches nor falls back, since every fallback allocates an output of the same size; set_*_enabled(True) clears a latch, and the process-local nature of latches is documented with its multi-rank consequence. Whole-module unpickling of a pre-fusion model works again, the fallback no longer graph-breaks under fullgraph, and an unapplicable activation raises instead of silently doing nothing. CPU 386 passed, GPU 267. Mutations: wiring 36/38 (two known-benign), kernel 35/45 with 9 pre-existing not-applicable, opt 48/50.
…ompiles activation_checkpointing with the compiled rung and a DCTensor input crashed: Dynamo hit a recompile limit reported as 8 while _raise_recompile_limit() had set 64, and setting 256 by hand did not help either. torch._dynamo.config's user overrides are ContextVar-backed, so they are per-thread. Non-reentrant checkpointing recomputes inside the backward pass, which runs on the autograd engine's device worker thread, and that thread reads the stock default. Instrumenting the limit check shows it directly: MainThread limit_seen=64 exceeded=0, worker thread limit_seen=8 exceeded=1. The compile that overflows is the one our setting cannot reach. Passing recompile_limit to torch.compile fixes it, because Dynamo applies that with config.patch() around the compile itself, on whichever thread compiles. What made it overflow is DistConv. TORCH_LOGS=recompiles shows the recompute's first miss is GLOBAL_STATE changed: torch_function -- DistConv's backward runs below __torch_function__, so the recompute misses all five forward entries and compiles a second parallel set. Five shapes times two torch-function states is ten, over the limit of eight. The count is bounded, so the fix is to make the limit real rather than to chase a guard: it now stabilizes at ten. FailOnRecompileLimitHit also escaped the ladder entirely, being the one Dynamo compile-time failure that derives from Exception rather than TorchDynamoException. It is named explicitly now. Catching it is not sufficient on its own, and this was measured rather than reasoned: with the exception merely caught, the real UNet died with CheckpointError and the synthetic case with a memory fault, because falling back mid-replay makes a checkpointed block's recompute disagree with its forward. So a proven rung now re-raises instead of degrading while a backward is in flight, detected with the graph-task id that checkpoint's own unpack hook uses. A first failure still degrades, which is where the contract about not killing a multi-node run actually lives. Two upstream bugs worth filing: the global recompile_limit does not apply to compiles triggered from the autograd worker thread, and its warning prints the limit it enforced while the config reads the value we set; and FailOnRecompileLimitHit sits outside the TorchDynamoException root that handlers are written against. CPU 391 passed, GPU 270. Wiring mutations 37/38, the survivor equivalent. Interleaved A/B on the Triton hot path: 0.1651 vs 0.1651 ms.
The premise this started from was wrong in two ways, both worth recording. Most of the profile's 4.77 ms aten::cat is not the skip concatenation at all. forward_halo_exchange ends in a cat and distconv_forward calls it once per spatial dim, so every 3x3x3 convolution pays three full-tensor concatenations to materialize a (D+2, H+2, W+2) copy -- including at dc_num_shards=1, where it is an open-coded padding=1. At the largest decoder block that is 1.90 ms of halo cat against 0.62 ms of skip cat. And ATen's cat is not making the mistake Inductor made with GroupNorm: on the forward it runs at 100.5 / 104.6 / 107.5% of the 3.35 TB/s streaming roofline at the three large decoder shapes. A Triton kernel with the same dtype behaviour measures -0.03 ms, i.e. nothing. What is actually wrong is the dtype and the backward. Under autocast the skip arrives fp32 (GroupNorm's fp32 policy) and the upsampled half bf16; cat carries the promote policy, so it widens bf16 to fp32, concatenates at fp32, and the convolution narrows it back -- three full-resolution passes to deliver one. And cat's backward hands out views that its consumers then force contiguous, at 51-63% of roofline. skip_concat does it in one channels-last pass at the consumer's dtype, with a split backward that addresses its cotangent by stride instead of calling contiguous() on it. That last part is the whole difference: under DCTensor the cotangent is a narrowed view of the halo-padded tensor, and the obvious kernel lost to plain torch (-1.06 vs -1.85 ms) until it stopped materializing it. Whole scale-7 step 92.783 -> 91.605 ms (1.3%), peak 7.215 -> 6.715 GiB. The four decoder blocks in isolation 53.053 -> 50.732 ms (4.4%); the largest block's skip path 2.43 -> 0.91 ms, with convolution_backward flat at 17.43/17.40 so nothing moved into the convolutions. Every conv output, block output and input gradient is bitwise identical and so is the whole-model loss. Measured and rejected: splitting the first convolution by input channel, which is algebraically equivalent and 5.5% slower at every block -- two half-channel convolutions plus an add cost more than one convolution plus a concatenation. CPU 408 passed, GPU 339.
Of the 1.18 ms ff2813c took off the scale-7 step, 1.09 ms and all 0.50 GiB of the memory came from concatenating at the dtype the following convolution reads instead of at torch.cat's promoted one. The Triton kernel was worth the remaining 0.08 ms -- under 0.1% of the step. That does not pay for 824 lines of hand-written kernel plus its 69 tests in a benchmark whose value depends on other people trusting it, particularly so soon after a review round found nine defects in the first one. The dtype computation moves into unet_parts as _consumer_dtype/_skip_concat, which is what the kernel's own fallback path already was. Behaviour outside autocast is unchanged, so eval, inference_mode and pure-fp32 runs still get torch.cat's ordinary promotion. What the kernel measured is kept in the Up docstring rather than the code: cat's backward is a narrowed view that consumers force contiguous, at 51-63% of this device's streaming roofline against the kernel's 90-103%. If the skip path is ever worth revisiting, that is where the remaining headroom is. CPU 397 passed, GPU 270.
Collaborator
Author
|
Closing in favor of a broader Triton PR. |
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.
(This will be rebased onto main once #92 lands. If you want to look at the new commits, they start from
199251867f2dc8242190d0e5240769668b03bb34.)DCTensorso we can actually compile GroupNorm. Compiled GroupNorm gets up to ~10x improvements over the eager-mode GroupNorm, depending on input shapes.