Enable and optimize prof-trainer on AMD Instinct MI355X (ROCm) - #4
Open
jbelof wants to merge 2 commits into
Open
Enable and optimize prof-trainer on AMD Instinct MI355X (ROCm)#4jbelof wants to merge 2 commits into
jbelof wants to merge 2 commits into
Conversation
This change makes prof-trainer run correctly on the AMD Instinct MI355X
(gfx950 / ROCm 7.2 / OpenMPI 5.0.x) and adds a set of performance
improvements, measured on the analytical Fourier-modes example dataset
(14,641 images, 512x512, fp32). It also adds env-gated experimental
toggles for further tuning. All measurements below are single-node.
=====================================================================
Correctness / enabling fixes (always on)
=====================================================================
1) Initialize the HIP context before mpi4py imports MPI.
The module previously imported `from mpi4py import MPI` before torch
("MUST BE IMPORTED BEFORE torch!"). On the ROCm + OpenMPI stack this
is fatal: once MPI_Init has run, the HIP runtime can no longer
initialize, so torch.cuda.is_available() returns False and
init_process_group("nccl") aborts with
"ProcessGroupNCCL is only supported with GPUs, no GPUs found!".
We now import torch first and force torch.cuda.init() before importing
mpi4py. Without this fix prof-trainer does not start at all on this
platform (single- or multi-GPU).
2) Make MASTER_PORT overridable.
MASTER_PORT was hardcoded to "23456", which prevents running more than
one single-GPU job per node. Changed to os.environ.setdefault(...) so
the default is unchanged for a normal (multi-rank) job, but independent
jobs can be placed on different GPUs/ports.
3) Guard the ZeRO-only optimizer call.
optimizer.consolidate_state_dict() is only defined on
ZeroRedundancyOptimizer; it is now called via hasattr(...) so an
alternative optimizer (see PROF_PLAIN_ADAM below) can be used without
crashing at checkpoint time.
=====================================================================
Performance optimization (always on)
=====================================================================
4) Accumulate per-step training metrics on-GPU, sync once per epoch.
The training loop previously called consolidated_loss() and then
.item()/.cpu() on four quantities (l1-sum, l2-sum, linf,
sum-per-channel) every step. Each of those is a device->host sync that
stalls the pipeline. We now accumulate them in on-device tensors and
copy to the host exactly once per epoch, right after the step loop.
TensorBoard/logging output is numerically identical to before.
Measured: ~1 s/epoch (~3.5%) at batch 255 on one MI355X.
=====================================================================
Experimental toggles (env-gated, default OFF -> stock behavior)
=====================================================================
- PROF_DROP_LAST=1 : drop_last on the train DataLoader. Avoids a
second MIOpen kernel-compilation pass for the
odd final (partial) batch. Critical for a cold
first run (see below); ~0.7% data cost.
- PROF_AMP=bf16|fp16 : wrap forward+loss in torch.autocast. bf16 gives
a large speedup on this hardware (see below).
- PROF_CHANNELS_LAST=1 : run the model/inputs in channels_last (NHWC).
- PROF_CUDNN_BENCHMARK=1: set torch.backends.cudnn.benchmark. NOTE: on
gfx950 this triggers an extremely long MIOpen
exhaustive find; not recommended.
- PROF_PLAIN_ADAM=1 : use torch.optim.Adam instead of
ZeroRedundancyOptimizer (no benefit on 1 GPU;
provided for comparison).
Also adds data/analytical-example/gen_nopng_mpi.py: an MPI dataset
generator that writes HDF5 identical to fouriermodes_study_mpi.py but
skips the matplotlib PNG render (PNGs are not used for training),
generating the full 14,641-image set in ~30 s across 64 ranks.
=====================================================================
Measured results on a single AMD Instinct MI355X (gfx950, ROCm 7.2)
Analytical example, 512x512, fp32, per-GPU batch 255 (see note), L1 loss,
max_feature 1024, min_feature 128, y/x_kernel 4. Model: 61.4M params.
=====================================================================
Steady-state, single GPU, fp32:
- Original per-step-sync loop, default MIOpen find : ~28.5 s/epoch
- + MIOPEN_FIND_MODE=2 (FAST) : ~29.5 s/epoch
- + once-per-epoch metric sync : ~28.4 s/epoch
- + torch.compile (--compile) : ~27.3 s/epoch
(recommended fp32 stack; ~530 img/s)
Cold first-run (empty MIOpen cache), single GPU, fp32, ONE epoch:
- default MIOpen find mode : ~231 s, and can exceed 15 min
when the odd final batch forces a
second on-GPU kernel search
- MIOPEN_FIND_MODE=2 + PROF_DROP_LAST=1 : ~33 s (~7x faster)
The dominant cold-start cost is MIOpen runtime kernel compilation, not
the ~28 s of training math; MIOPEN_FIND_MODE=2 avoids benchmarking many
candidate kernels on-GPU on the first pass.
Precision (single GPU, optional, PROF_AMP=bf16):
- fp32 : ~28.4 s/epoch
- bf16 : ~10.1 s/epoch (2.8x faster), ~54 GB vs ~74 GB VRAM,
training-loss trajectory unchanged.
Multi-GPU strong scaling (recommended fp32 stack, fixed 14,641-image
dataset split by DistributedSampler; steady-state):
- 1 GPU : 27.45 s/epoch, ~530 img/s, 1.00x
- 4 GPU : 6.99 s/epoch, ~2043 img/s, 3.86x (96.5% efficiency)
- 8 GPU : 3.76 s/epoch, ~3797 img/s, 7.17x (89.6% efficiency)
Near-linear; the compute-heavy generator hides the ~245 MB fp32
gradient all-reduce. The 8-GPU dip is a strong-scaling artifact
(only 7 steps/rank, so fixed per-epoch overhead dominates a short
epoch), not a communication wall.
Recommended launch (single GPU, cold-start friendly, multi-epoch):
MIOPEN_FIND_MODE=2 PROF_DROP_LAST=1 \
MIOPEN_USER_DB_PATH=/tmp/miopen_cache MIOPEN_CUSTOM_CACHE_DIR=/tmp/miopen_cache \
mpirun -n 1 --bind-to none -x ROCR_VISIBLE_DEVICES=0 \
-x MIOPEN_FIND_MODE -x PROF_DROP_LAST \
-x MIOPEN_USER_DB_PATH -x MIOPEN_CUSTOM_CACHE_DIR \
prof-trainer --batch_size 255 --num_epochs 100 --compile ...
Recommended launch (N GPUs): mpirun -n N --bind-to none (drop the
ROCR_VISIBLE_DEVICES mask; each rank selects its GPU via local_rank).
Note on batch size: the MIOpen spatial-BatchNorm kernel on gfx950
computes N*C*H*W in int32. For this model the largest BN layer is
[N,128,256,256] (per-sample 8,388,608), so per-GPU batch must satisfy
N*8,388,608 < 2^31, i.e. batch_size <= 255. Larger per-GPU batches abort
with miopenStatusUnknownError; global batch scales as 255*num_gpus.
=====================================================================
Portability to other ROCm / CDNA hardware (e.g. MI300A)
=====================================================================
These results were measured on the MI355X (CDNA4 / gfx950), but most of
the improvements are stack-level rather than gfx950-specific and are
expected to carry over to other ROCm parts such as the MI300A
(CDNA3 / gfx942), though absolute numbers will differ:
- Very likely to transfer (software/stack, not arch-specific): the
HIP-before-MPI init fix, the MASTER_PORT override, the
consolidate_state_dict guard, the once-per-epoch metric sync, and the
cold-start MIOPEN_FIND_MODE=2 + PROF_DROP_LAST behavior (MIOpen
find/compile cost is common across CDNA GPUs). torch.compile is a
framework feature and is architecture-independent.
- Directionally applicable, magnitude will differ: bf16 autocast
(CDNA3 also has bf16 matrix throughput well above fp32, so a large
speedup is expected, but not necessarily 2.8x) and multi-GPU scaling
(the "compute hides the all-reduce" argument is general; note MI300A
nodes are typically 4 APUs with unified memory, so the topology and the
8-GPU point specifically differ).
- NOT assumed for MI300A: the batch_size <= 255 int32 BatchNorm cap was
observed with a specific gfx950 MIOpen batchnorm kernel; MIOpen may
select a different solver on gfx942, so that limit may or may not
appear there and should be re-verified.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add docs/mi355x_partition_modes.md documenting how the MI355X compute/memory partition modes (SPX/DPX/QPX/CPX x NPS1/NPS2, selected per Slurm job via --comment) affect prof-trainer throughput on the analytical example with the optimized fp32 stack. Key finding: - CPX/NPS1 is a good optimization for a SINGLE GPU: ~+10% throughput vs the default SPX (8 concurrent smaller-batch partitions fill one physical GPU better than a single large-batch stream for this ConvTranspose model). - CPX is NOT suitable for MULTI-GPU training. Its per-GPU advantage reverses at 4 GPUs (32-way DDP: 1418 vs SPX 2043 img/s, -31%) and it segfaults at 8 GPUs (64-way, at the first step, independent of batch size or --compile). Use SPX (whole GPUs) for any DDP run; SPX scales near-linearly (3.86x/7.17x at 4/8). - NPS1 vs NPS2 is within noise for this workload. Recommendation: CPX/NPS1 for single-GPU jobs, SPX/NPS1 for multi-GPU. Also adds the page to the mkdocs nav under a new "Performance" section. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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.
Summary
Makes
prof-trainerrun on the AMD Instinct MI355X (gfx950 / ROCm 7.2 /OpenMPI 5) — it does not start out of the box on this stack — and adds measured
performance improvements plus MI355X partition-mode guidance. Validated on the
analytical Fourier-modes example.
cc @cjekel @cssherman
Changes
Correctness / enabling (always on)
MPI_Init(otherwisetorch.cudasees no GPUs and
init_process_group("nccl")aborts).MASTER_PORToverridable via env (was hardcoded to 23456).consolidate_state_dict()so alternative optimizers work.Performance (always on)
~4 device→host stalls per step; TensorBoard output is numerically identical).
Experimental toggles (env-gated, default off):
PROF_DROP_LAST,PROF_AMP=bf16|fp16,PROF_CHANNELS_LAST,PROF_CUDNN_BENCHMARK,PROF_PLAIN_ADAM.Tooling / docs
data/analytical-example/gen_nopng_mpi.py: fast MPI dataset generator (HDF5identical to the existing script, no PNG render; ~30 s for 14,641 images).
docs/mi355x_partition_modes.md: SPX/DPX/QPX/CPX × NPS1/NPS2 findings.Measured on a single MI355X (analytical example, 512×512, fp32)
FIND_MODE=2+drop_lastMulti-GPU (fp32, strong scaling): 1/4/8 GPU = 1.00× / 3.86× / 7.17× (96% / 90% eff).
Partition modes (single GPU): CPX/NPS1 ≈ +10% vs SPX — but CPX does not
scale (−31% at 4 GPUs, segfaults at 8 GPUs), so use SPX for multi-GPU. Details
in the new doc.
Portability
Most changes are ROCm-stack-level and should carry to other CDNA parts (e.g.
MI300A / gfx942); absolute numbers will differ, and the batch ≤ 255
int32-BatchNorm cap is gfx950-specific and should be re-verified elsewhere.
Testing
Built and run on 1 node of 8× MI355X (ROCm 7.2, PyTorch 2.11+rocm7.2, OpenMPI
5.0). Full pipeline validated (MPI init, RCCL, data loading, checkpointing,
TensorBoard) at 1/4/8 GPU; metric output confirmed identical to the previous
per-step path.
🤖 Generated with Claude Code