[MLPerf 6.1][DLRMv4] Adding DLRMv4 HSTU - #889
Open
chriscai-amd wants to merge 155 commits into
Open
Conversation
…7e51c)
Vendored snapshot of chriscai-amd/generative-recommenders branch chcai/dlrmv4
(HEAD d97e51c) as a sibling of recommendation_v2/torchrec_dlrm. The Python
package generative_recommenders keeps its original name so all imports work
unchanged from the new location.
- recommendation_v4/generative_recommenders/: dlrm_v3, modules, ops, research, tests
- recommendation_v4/configs/: research HSTU gins
- recommendation_v4/scripts/launch_smoke_8gpu.sh: sanitized 8-GPU yambda-5b launcher
(resolves package root from script path; AMD env defaults; pip_local override)
- recommendation_v4/{setup.py,requirements.txt,main.py,...}: upstream entry points
- .gitmodules: cutlass registered at parent repo level
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Four fixes unlocking the HSTU_HAMMER_KERNEL=TRITON path on MI350X: 1. triton_hstu_attention.py _should_enable_tma(): add HIP early-out. torch.cuda.get_device_capability() on gfx950 returns (9, 5) which would pass the major==9 Hopper check and trick the kernel into the TMA path, producing kernels that don't compile on ROCm. 2. triton_hstu_attention.py _get_fw_configs(): hoist the USE_TLX/NUM_BUFFERS/ NUM_MMA_WARPS_PER_GROUP/NUM_MMA_GROUPS defaults loop out of the CUDA-only else: branch. The _hstu_attn_fwd signature requires these constexprs regardless of backend; missing them on HIP triggered TypeError: dynamic_func() missing N required positional arguments at autotune. Also gate the H100 TLX configs append on `not torch.version.hip`. 3. triton_jagged_tensors.py concat/split dispatch: route AMD/ROCm through *_2D_jagged_multirow instead of the basic _concat_2D_jagged / _split_2D_jagged kernels. The basic kernels fail PassManager::run at make_ttgir (TritonAMDGPUCanonicalizePointers pass) on ROCm; multirow compiles fine. NVIDIA non-Blackwell paths (H100/A100) are unchanged. 4. triton_jagged_tensors.py _Concat2DJaggedFunction.backward: replace the raw _split_2D_jagged[grid] call with _triton_split_2D_jagged_internal so the backward pass benefits from the same AMD multirow routing as the forward. Verified end-to-end on 8x MI350X: yambda-5b bs=32 seq=4k at 782 global_sps vs PYTORCH backend 547 sps -- 1.43x throughput, 75% peak VRAM vs 92%. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The attribute is absent in some Triton builds (e.g. nvcr.io/nvidia/pytorch:26.01-py3), causing import-time AttributeError before any training step runs. Use getattr with a False default so _use_meta_ws() gracefully reports disabled on those builds.
Three small changes so you can sweep model size and per-sample sequence
length from a gin file without editing configs.py.
configs.py:
- get_hstu_configs is now @gin.configurable. Accepts optional overrides
for max_seq_len, max_num_candidates, hstu_embedding_table_dim,
hstu_transducer_embedding_dim, hstu_num_heads, hstu_attn_num_layers,
hstu_attn_linear_dim, hstu_attn_qk_dim, hstu_input_dropout_ratio,
hstu_linear_dropout_rate. Per-dataset defaults still apply unless
explicitly overridden in gin.
- get_embedding_table_config is now @gin.configurable with an
embedding_dim override that uniformly sets the dim for all tables
of the chosen dataset.
- Drop the YAMBDA_EMBEDDING_DIM constant (was a duplicate of
HSTU_EMBEDDING_DIM=512). Yambda branch now uses HSTU_EMBEDDING_DIM
directly. Add a comment noting the model+table dim must stay aligned
when overriding either via gin.
utils.py:
- get_dataset accepts an optional history_length kwarg that wins over
the yambda dataset's hardcoded default of 4096. Caches are still
keyed on disk under hstu_cache_L<N>/ so switching L between previously
built values is free.
train/gin/yambda_5b.gin:
- Pin history_length=2048 and max_seq_len=2048 for the seq-2k smoke
config. Both lines have inline comments explaining the +9 overhead
(uid + 7 cross + 1 candidate) so total per-sample seq is ~2046,
within the 2048 budget.
Verified: default codepath unchanged, gin overrides apply consistently
to both get_hstu_configs (model) and get_embedding_table_config (tables).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
make_optimizer_and_shard now accepts hbm_cap_gb (default 260, the MI350X value) via @gin.configurable. The yambda gin pins the same default so sweeps just change the number in the gin file instead of editing utils.py. ddr_cap dropped from 32 GiB to 0: with all 11 yambda 5b embedding tables fitting on 8x MI350X HBM, allowing host DRAM offload only invites the planner to pick slower per-lookup-PCIe-traffic plans. Verified gin binding flows through to the Topology: a probe with hbm_cap_gb=100 produced Topology(hbm_cap=107374182400) and the planner correctly raised insufficient-storage error at that tightness. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
preprocess_public_data.py:
- Add DLRMYambdaProcessor: downloads Yambda multi_event + catalog
metadata from the yandex/yambda HuggingFace repo, then runs a
temporal split (300 train days / 30 min gap / 1 test day),
builds per-user sessions (1800s inactivity threshold), and
writes the layout DLRMv3YambdaDataset expects:
<data-path>/raw/<size>/multi_event.parquet
<data-path>/shared_metadata/{artist,album,embeddings}.parquet
<data-path>/processed_<size>/{train_sessions,test_events,
session_index}.parquet
<data-path>/processed_<size>/item_popularity.npy
<data-path>/processed_<size>/split_meta.json
- 5b variant uses chunked polars load (10M rows/chunk) to keep
peak RAM under control (single-shot read of the 50 GB parquet
OOMs ~150 GB systems).
- SUPPORTED_DATASETS adds yambda-50m, yambda-500m, yambda-5b.
- main() takes --data-path for custom output root.
- Verified end-to-end: 50m run completes in ~2 min, 5b in ~53 min
(download dominates), output is byte-compatible with the dataset
cache builder; TRITON training reaches steady state on the
fresh data at 2050 sps.
utils.py:
- Add env_path(key, default) @gin.configurable helper. Used as a
gin macro so any string-valued binding can be overridden by an
env var without editing the gin file.
train/gin/yambda_5b.gin:
- Declare DATA_PATH = @env_path() macro with key="DLRM_DATA_PATH"
and default="/apps/chcai/dlrm_data". Both new_path_prefix
bindings (make_train_test_dataloaders and get_dataset) now
consume %DATA_PATH. Setting DLRM_DATA_PATH=/some/path at run
time redirects the dataset without a gin edit.
datasets/yambda.py:
- Strip stale references to upstream-internal preprocessing in
docstrings/comments; point at preprocess_public_data.py instead.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Every rank's first CUDA context was landing on GPU 0 (the default device), so NCCL bound its communicators there before set_device switched to the correct GPU. This leaked allocations on GPU 0 across all 8 ranks and caused spurious OOMs during embedding-table init at high HBM caps. Moving set_device above init_process_group and passing device_id ensures each rank's NCCL state is created on its own GPU.
dlrm_v3/utils.py:
- Replace the hardcoded manifold:// URL in _on_trace_ready_fn with a
local trace_dir (default /tmp/dlrm_v3_traces). Filename now follows
trace_step{step}_rank{rank}.json so per-rank captures don't collide.
- Add _multi_window_schedule helper: a torch.profiler schedule that
fires around each step in trace_steps=[...] (warmup before, active
after, RECORD_AND_SAVE at the last active step). Lets one run
capture multiple windows (e.g. early-step + steady-state) without
re-running.
- Make Profiler @gin.configurable. New knobs: trace_dir, trace_steps,
wait, warmup, repeat, record_shapes, profile_memory, with_stack,
with_flops, with_modules. Defaults preserve the prior single-window
behavior (wait=10, warmup=20, active=50, repeat=1) so existing
callers are unaffected.
- Add run_results_dir(run_name) gin macro: resolves to
<recommendation_v4>/results/<run_name>/. Used as the canonical
output prefix for traces (and any future per-run artifacts).
recommendation_v4/ is bind-mounted into the training container, so
files written through this helper persist on the host.
train/gin/yambda_5b.gin:
- Wire RUN_NAME env override -> run_results_dir(run_name=%RUN_NAME)
-> Profiler.trace_dir. Sets trace_steps=[52], warmup=5, active=5
(capture the 5-step window 52-56 on every rank).
- Toggle train_eval_loop.output_trace = True so the profiler actually
instantiates.
.gitignore:
- Add results/ alongside the existing tmp/exps/ckpts/ runtime
directories so per-run trace dumps don't show up in git status.
Verified: 8x MI350X TRITON yambda-5b run at bs=32 seq=2k drops
8 well-formed trace_step62_rank{0..7}.json files (~37 MB each) into
recommendation_v4/results/default/; visible on the host immediately.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…mes, trim_warmup
dlrm_v3/utils.py
* Add run_results_dir(run_name) gin macro (resolves to
<recommendation_v4>/results/<run_name>/) so trace artifacts persist on
the host via the bind-mount.
* Add _trim_warmup_from_trace post-processor: dedupes ProfilerStep spans
by name first, then keeps only the last N unique steps' worth of
events. Drops WARMUP-phase events that torch.profiler otherwise
includes in the chrome trace.
* Add trim_warmup kwarg (default True) on Profiler; auto-invokes the
trimmer with N=active so the exported file matches the user-requested
active window.
* Filename now uses trace_steps[i] (the user-requested step) as the
{step} label when multi-window mode is in use, instead of
torch.profiler's internal step_num (which is off by ~warmup+active
from the schedule trigger and confused everyone).
train/utils.py
* Drop hardcoded `active=10` from the four `Profiler(rank, active=10)`
call sites in train_loop / train_eval_loop. Positional args block
gin overrides; once removed, Profiler.active in gin (default 50) and
user gin bindings actually take effect.
train/gin/yambda_5b.gin
* Fix env_path scoping collision: both DATA_PATH and RUN_NAME used the
unscoped @env_path() configurable, which made the second binding's
`env_path.key = "RUN_NAME"` overwrite the first's
`env_path.key = "DLRM_DATA_PATH"`. Both names then resolved via the
same env var (whichever was last), pointing DATA_PATH at trace_run2/
and breaking dataset loads.
Fixed by giving each call site its own scope: @data/env_path() and
@run/env_path(), each with independent .key/.default bindings.
* Set Profiler.trace_steps=[52], warmup=1, active=5; let trim_warmup
default to True so the exported trace contains exactly 5 active
ProfilerStep events.
Verified end-to-end:
- Run with RUN_NAME=trace_run2 writes results/trace_run2/trace_step52_
rank{0..7}.json (~19 MB each), step labels match trace_steps gin.
- Triton cache persisted across runs: cold start ~6 min -> warm start
~2 min for autotune-to-first-step.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2048 was chosen for "round number near max_seq_len" but it slightly
overflows the per-sample budget: 3 * (2048//3) + 9 = 2055 > 2048, so
the dataset truncates ~7 UIH events to fit. 2039 makes the math exact
(3 * 679 + 9 = 2046 ≤ 2048) so no truncation.
Comment block expanded to document:
- The 3-pool gather semantic (L//3 events per pool, interleaved
chronologically).
- The like-pool under-fill observation: like events are only 1.9%
of yambda corpus and max user lifetime is ~28k events, so the
like pool fills to ~105 events per anchor on average (not 679).
TRITON's jagged attention skips the unfilled slots, so under-fill
costs sequence budget but not GPU compute.
No code change. Cache for L=2039 already built and reused.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…ol gather Documents the fork's scope (yambda-5b on HSTU dlrm_v3 path), per-pool gather strategy with effective fill table, and dataset statistics. Sections indexed 1–5 for navigation. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Adds env_int gin macro (companion to env_path) and wires make_optimizer_and_shard.hbm_cap_gb through it so the per-rank HBM ceiling can be tuned without editing the gin file. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Document the container image, dependency versions (native NGC torch 2.10, triton 3.6, source-built fbgemm_gpu, torchrec 1.4.0, polars-u64-idx), gin training configuration, and env vars needed to reproduce the 8x B200 run.
Adds three knobs, all driven from the gin file: - make_model.bf16_training: enable bf16 autocast for the DlrmHSTU model. - env_int macro: lets numeric gin values come from env vars (used by the existing hbm_cap_gb binding). - apply_env_bootstrap.TRITON_FULL_AUTOTUNE: when False (default), three layer-norm/jagged triton kernels are pinned to a single Config so cold starts land at the same steady-state deterministically. When True, the full autotune search runs again — use this when changing shape, GPU, or triton/torch version, then re-pin from the discovered winners. train_ranker._main_func now parses gin in two phases (skip_unknown=True early, full pass after the heavy imports) so the bootstrap env var is set BEFORE the triton kernel modules evaluate their @triton.autotune decorators at module load time. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Mirrors the B200 layout with MI350X (gfx950, ROCm 7.2.1) specifics: container image (rocm/primus:v26.3), fbgemm_gpu rebuild requirement (HEAD nightly_rocm-2026.6.1 for ~30% step-time win over the shipped 2026.5.14), the gin-driven TRITON_FULL_AUTOTUNE knob, and the measured perf ladder from fp32/PYTORCH baseline (~28 d/epoch) down to the pinned bf16/TRITON fast equilibrium (~7.6 d/epoch). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Merges per-rank chrome traces (results/<run>/trace_step{N}_rank{R}.json)
into a single Perfetto-loadable file, remapping pid/flow ids so
cross-rank events land on distinct tracks instead of collapsing onto one.
Used to produce the bf16 + pinned-autotune step-52 trace
(results/verify_rename/trace_step52.json.gz).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Refresh the B200 dependency versions to the latest validated stack (torch 2.12.0a0 / CUDA 13.2, fbgemm_gpu built for sm_100+CUDA 13.2, CUPTI 13.2), note 26.01 as an equivalent alternative, and record the TRITON_FULL_AUTOTUNE=True setting for B200.
…iver) Point fbgemm at the latest validated source commit (10b77573, 2026-06-01), record the tested torchrec 1.7.0.dev nightly (1.4.0 stable fallback), clarify the fbgemm wheel version string is the build date, and correct the host/forward-compat driver CUDA versions (13.0 host / 595.58.03 compat).
After upgrading to torch 2.12 / torchrec 1.7 (B200-aligned), the pinned configs from the torch 2.10 stack stopped landing on the fast equilibrium because the torchrec 1.7 code path invokes these kernels at different shape keys. Re-captured winners via a fresh autotune run and updated the pin sites: - _weighted_layer_norm_bwd_dx: BLOCK_N 8 -> 1 (num_warps 1 unchanged) - split_2D_jagged_multirow: BLOCK_N 1 / num_warps 2 -> BLOCK_N 8 / num_warps 1 - _layer_norm_bwd_dwdb: BLOCK_N 128, num_warps 8 (unchanged - same winner on both stacks) Verified: 3 consecutive checkpoints (steps 151/201/251) at 52.75-53.36 ms deterministic on the new stack. Same equilibrium band as the torch 2.10 stack (51.5-53.0 ms). Also adds a Stack B section to docs/training_recipe.md (MI350X) documenting the torch 2.12 swap recipe (torch + torchvision + torchaudio + fbgemm rebuild + torchrec git tag) so the MI350X recipe is dependency-aligned with the B200 path. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Bumps the Stack B (torch 2.12 / torchrec 1.7) section to: - fbgemm commit 10b77573 (same SHA as the B200 path) instead of 1509423 (one cosmetic commit behind). Wheel rename 2026.6.1 -> 2026.6.2. - Note that Stack A and Stack B use different pinned triton configs (already merged) and explain why (torchrec 1.7 invokes the kernels at different shape keys). - Caveat: HSTU_HAMMER_KERNEL=PYTORCH fallback regresses to ~169 ms on Stack B (vs 107 ms on Stack A). TRITON is unaffected and remains the default; this only matters for PYTORCH-backend debugging. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Collapses the two-stack MI350X section into one canonical dependency table: torch 2.12 / torchrec 1.7 / fbgemm @ 10b77573 — the same SHAs as the B200 path. The image-native torch 2.10 / torchrec 1.4 / fbgemm 2026.5.14 path still works for development but the recipe doc now documents the validated production stack only. PYTORCH-backend caveat preserved. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Not relevant — TRITON is the documented default backend. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…ecipe Embedding sizes match the true entity counts in yambda-5b: item_id 9_390_000 -> 9_390_624 artist_id 1_290_000 -> 1_293_395 album_id 3_370_000 -> 3_367_692 uid 1_000_000 -> 1_000_001 This eliminates the recurring "EmbeddingBoundsCheck ... Setting idx to zero" warnings at training time. Gin default raised to batch_size=1024 / eval_batch_size=1024. Measured steady-state on the torch 2.12 + torchrec 1.7 + fbgemm HEAD stack with TRITON HSTU + pinned triton configs: ~635 ms/step, ~12.9K sps, ~2.92 days/epoch vs ~7.6 days at bs=32. bs=2048 is feasible but only +3% throughput at much higher autotune cost, so bs=1024 is the sweet spot. Triton autotune pin for _weighted_layer_norm_bwd_dx now ships TWO configs in the pinned list — BLOCK_N=1 (bs=32 winner) and BLOCK_N=8 (bs=1024 winner). Triton's autotune key=[BLOCK_D] dispatches the right one per shape in <5 sec on cold start (vs ~30 sec from the full pool). The other two pinned kernels (_layer_norm_bwd_dwdb, split_2D_jagged_multirow) have identical winners at bs=32 and bs=1024 so they stay single-config. Training-recipe doc drops the batch_size rows from both MI350X and B200 config tables — the recipe is intentionally batch-size-agnostic now that the pin set covers a range. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Enable the multi-row, separated-RNG _ln_mul_dropout path on AMD MI350 (gfx950), previously Blackwell-only. Batches rows per program and reuses a precomputed dropout mask in the backward instead of one-program-per-row fused RNG; +5.6% end-to-end (-> 14,222 global sps) at bs=1024 on yambda-5b. - ops/utils.py: add is_amd_mi350() + use_separated_rng_ln_mul_dropout() gate. - ops/triton/triton_hstu_linear.py: dispatch the fwd LN-dropout to the separated-RNG path via the new gate. - ops/triton/triton_hstu_attention.py: pin fast nonkdim:16 fwd/persistent/bwd configs via pinned_or_full (TRITON_FULL_AUTOTUNE=1 still bypasses). Multi-config lists with an inline "add a new batch size" guide. - scripts/launch_smoke_8gpu.sh: GPU clock sanity guard - log perf level + sclk, auto-restore 'auto' if a perf_determinism/manual/low lock is found (a half-clock lock uniformly slowed every Triton kernel ~1.9x and masked perf changes). - docs/perf_opt.md: document the LN-dropout fix and the clock-lock caveat. Co-authored-by: Cursor <cursoragent@cursor.com>
…ernel Add an opt-in TrainPipelineSparseDist path that overlaps the embedding input-distribution all-to-all with dense fwd/bwd. To make the embedding collection pipelineable, the merged sparse KJT is now pre-built in the dataloader (Samples.merged_sparse_features) and the model consumes it via a _pipeline_mode forward that takes the batch as a single arg, so TorchRec's tracer resolves the lookup input as a plain getattr off the batch. - dataset.py: Samples.merged_sparse_features + merge_uih_candidate_kjts, built in collate_fn; wired into to()/record_stream()/pin_memory(). - dlrm_hstu.py: _pipeline_mode flag; forward unpacks the batch and preprocess accepts the prebuilt merged KJT (falls back to building it when absent). - utils.py: _PipelineModelWrapper, build_train_pipeline, train_eval_loop use_pipeline branch + eval batch-arg; seed all RNGs in setup() for reproducible weight init. - gin/launch: make_model.hammer_kernel selects TRITON vs PYTORCH (env override still honored); launch script defers to the gin default. use_pipeline defaults to False. Validated on MI350/ROCm 8-GPU: embedding collection is pipelined (input-dist a2a moves to hidden); model quality and throughput match the sequential path (seeded A/B). The exposed embedding-output a2a still dominates the step, so throughput is unchanged — pipelining is quality- and perf-neutral here. Co-authored-by: Cursor <cursoragent@cursor.com>
Add a forward-in-time streaming path: slice the timeline into fixed-duration windows (default 1 day), train window T then eval window T+1, enforcing no future leakage (across-window + causal-history guarantees). Make it the default mode in launch_smoke_8gpu.sh. Window-reset overhead is hidden via a persistent worker pool + double buffering (next window's index mask and first-batch prefetch overlap compute on a background thread) and eval-window prefetch one window ahead, dropping train/eval first-batch waits to ~1-3ms with no steady-state regression. Window selection uses a lazily-built, mmap'd anchor-timestamp cache so the default non-streaming path is unaffected. Also harden trace export (best-effort: IO/permission failures warn instead of crashing training) now that streaming enables output_trace by default, and document the path + knobs in the README. Co-authored-by: Cursor <cursoragent@cursor.com>
save_dmp_checkpoint.path now resolves from $CKPT_PATH and defaults to empty, so checkpoints (a full DMP is ~100s of GB, and the streaming loop always saves the final window) are off unless explicitly enabled. Also drop the stale training-recipe sentence claiming native torch is kept — it contradicts the dependency table, which replaces torch and keeps only the image's triton. Co-authored-by: Cursor <cursoragent@cursor.com>
Add in-process trace postprocessing in the profiler on_trace_ready callback to fix two ROCm/roctracer rendering artifacts that make MI350X traces look wrong in Perfetto (the timing is correct, only the layout): - _normalize_profilerstep_layout: collapse the fragmented GPU-side ProfilerStep#N spans (roctracer splits a step across the HIP null + compute streams) into one full-width span per step on the busiest compute stream, matching the CUDA look. - _deoverlap_gpu_slices: pull back sub-us kernel end timestamps so back-to-back kernels don't touch/overlap; Perfetto otherwise nests the later (long) kernel inside the tiny epilogue and clips it to zero width, hiding kernels like _hstu_attn_bwd. Leaves a ~1ns gap (exact end==start is just as fatal as an overlap) and leaves real nesting untouched. Both passes are gated behind _is_rocm() (torch.version.hip) so they are complete no-ops on CUDA/B200, which don't have these artifacts. All best-effort: failures degrade to a warning and never crash training. Co-authored-by: Cursor <cursoragent@cursor.com>
Add _deoverlap_gpu_annotations to the trace-export postprocessing, the annotation-boundary analog of the kernel de-overlap. Kineto projects the forward/backward phase annotations (## user_forward ##, ## item_forward ##, ## stu_* ##, ...) onto the GPU stream as a chain of end-to-end siblings. The absolute step timestamps are ~5.4e12 us, where a float64's quantum is ~1 ns, so a sibling boundary that should be coincident lands a few ns off; when the earlier sibling ends at/after the next one's start, Perfetto nests and clips the next span to a sliver -- e.g. the 100+ ms ## user_forward ## vanishes on some ranks/steps purely by rounding luck. Since annotations form a real nesting hierarchy (user_forward contains the stu_* spans and their kernels), this walks the per-track slice stack and only snaps a slice back when the next slice extends beyond it (siblings, not parent/child), guarding against trimming into a span's own descendants. It also snaps kernel tails that straddle an annotation boundary. Gated by _is_rocm() (no-op on B200/CUDA) and best-effort like the other passes. Verified end-to-end on an 8-rank MI350X run: ## user_forward ## renders 40/40 (was 9/40), total clipped annotations 1352 -> ~5. Co-authored-by: Cursor <cursoragent@cursor.com>
Make streaming-train-eval crash-resumable and add general checkpoint
cadence controls:
- Atomic checkpoint saves (.tmp dir + rename), keep_last_n pruning, and
swap-aside .old overwrite so a save can safely replace an existing
train_ts dir; stale .tmp/.old swept on the next save.
- Per-rank RNG snapshot/restore for bit-equal dropout replay on resume;
auto-latest-subdir resolution + (train_ts, batch_idx_in_window) resume
hint so a run re-enters a partial window and skips already-trained
batches exact-once.
- Three independent in-window checkpoint cadences via a pure, testable
decision helper: per-window batch count, monotonic global step
(e.g. every 1000 steps), and wall-clock interval (e.g. hourly,
rank-0-decided + broadcast to keep the save barrier in lockstep).
- gin/env bindings for all cadences + a test-only die_at_step hook.
Tests: checkpoint_cadence_test.py (cadence precedence/triggers) and an
end-to-end baseline/interrupt/resume harness (streaming_resume_test.{sh,py})
that gates on functional invariants (RNG restored, correct resumed step,
atomic save, keep_last_n) plus a loose trajectory-closeness bound.
Co-authored-by: Cursor <cursoragent@cursor.com>
Add $SKIP_EVAL (streaming_train_eval_loop.skip_eval, default 0 = OFF) to suppress
the first N scheduled full-holdout eval passes, cutting eval overhead early in a
run when the model is far from the AUC target and those points carry little
signal.
Works for both cadences, keyed off the same grid each already uses so the skip
set is resume-safe (derived from checkpoint-restored global_step / absolute ts,
never a per-process counter):
* data-fraction: skip while ordinal k = global_step / eval_interval_steps <= N;
eval begins at the (N+1)th data-pct point.
* per-window: skip while (train_ts - eval_anchor_ts)/eval_every_n_windows < N.
The end-of-run FINAL eval is never skipped; early-stop / MLPerf EVAL_ACCURACY
still fire once eval starts. Wired through launch_slurm.sh's worker passthrough.
Verified e2e (per-window, START_TS=150, NUM_TRAIN_TS=4, EVAL_EVERY_N_WINDOWS=1):
SKIP_EVAL=0 eval'd ts=150,151,152,153,final; SKIP_EVAL=2 eval'd ts=152,153,final
(first two points correctly dropped, final always run).
Co-authored-by: Cursor <cursoragent@cursor.com>
Add 10 more per-seed MLPerf logs to each gbs bucket (8192/16384/32768), bringing each to 20 runs. All 60 logs reach >=0.75 eval AUC; the three short single-node (gbs_8192) runs that stalled below 0.75 were replaced with full 8h runs. Co-authored-by: Cursor <cursoragent@cursor.com>
…g/dlrmv4-rcp-logs
…ip_eval=10) Update yambda_5b.gin env-default values so the standard config matches the convergence/RCP batch: DENSE_LR/SPARSE_LR 1e-7->1e-6, LR_WARMUP_STEPS 0->24000, EVAL_EVERY_DATA_PCT 0.005->0.001, SKIP_EVAL 0->10. Env overrides still take precedence, so existing launch scripts are unaffected. Co-authored-by: Cursor <cursoragent@cursor.com>
The 60 RCP reference logs are already tracked; remove the !*.log negation file now that the log set is final. Co-authored-by: Cursor <cursoragent@cursor.com>
dlrmv4: add RCP reference logs by global batch size
Summarize the merged rcp_logs sweep (3 global batch sizes x 20 seeds), add samples- and wall-clock-to-converge at AUC>=0.75, switch the eval target to 0.75 throughout, clarify the fixed-start/warmup LR schedule and 0.1% eval frequency, and add a 10.1 tunable-hyperparameters subsection. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
chriscai-amd
changed the title
chriscai-amd
marked this pull request as ready for review
July 30, 2026 17:30
The 60 RCP convergence logs recorded submission_platform=MI355X from the mlperf_logging_utils default, which disagreed with the MI350 platform name documented in README.MD section 10. Align all logs on MI350 (one line each, JSON re-validated across all 62,580 MLLOG records). Also remove the rcp/ placeholder directory: it only held a "not yet generated" README that rcp_logs/ has superseded, and nothing referenced it. Co-authored-by: Cursor <cursoragent@cursor.com>
….1%) The RCP logs were generated at EVAL_EVERY_DATA_PCT=0.001, but the launcher defaults and the README/gin comments still said 0.005, so a reference run launched without an explicit override would have evaluated 5x less often than the RCPs it is compared against. Co-authored-by: Cursor <cursoragent@cursor.com>
…POCH_PCT Skipping a fixed COUNT of early eval passes does not transfer across global batch sizes: convergence arrives later in epoch terms as GBS grows, so one count is either wasteful at small GBS or overshoots the earliest crossing at large GBS. Express the warm-up skip as a fraction of the epoch instead, which is GBS-independent in meaning and reads directly off the RCP epoch-to-converge curve. The fraction resolves once at startup into a global train-step threshold, keyed off the monotonic checkpoint-restored global_step rather than a per-rank sample tally, so the skipped set is identical across resume and every rank reaches the same verdict (a per-rank disagreement would deadlock the eval collective). Both the data-fraction and per-window cadences consult it; the end-of-run final eval is never skipped. Defaults derived from the 20-seed RCP sweep by fitting the earliest observed 0.75 crossings against GBS (README 10.2); 0.027 at GBS 8192. Cuts eval from ~29% of end-to-end time to ~5%. Also log the per-window skip in the double-buffer path, which previously skipped silently under the default DOUBLE_BUFFER=1. Co-authored-by: Cursor <cursoragent@cursor.com>
Reorganize the README onto the six-section spine of benchmark_readme_template.md (Problem / Directions / Dataset+Environment / Model / Quality / Approx runtime), following the text_to_image precedent, so reviewers can find the standard information where they expect it. Existing content is re-homed rather than rewritten; MLPerf-specific material (hyperparameters, RCP sweep, eval-schedule derivation, streaming knobs, checkpointing) is folded in as subsections. Fills four template requirements the README did not previously cover: dataset Publication/Attribution (arXiv:2505.22238 - the repo had no Yambda citation), weight and bias initialization, loss function as its own subsection, and evaluation thoroughness (eval_samples=2,058,204, no subsampling). Adds section 6 with the RCP wall-clock ranges on MI350X/bf16, and expands the model section into an explicit ordered layer list. Also makes the eval-cadence arithmetic explicit: 0.1% of the epoch is 2,290,835 samples but rounds up to a whole number of global batches, landing on 2,293,760 at all three RCP batch sizes. Co-authored-by: Cursor <cursoragent@cursor.com>
The hyperparameter table presented every parameter as tunable with a per-row "tuning rule", contradicting the submission rules in section 5 that allow only global batch size and target learning rate to be changed. Drop that framing and keep the section purely descriptive; tunability is stated once, in section 5. Trim the table to the main parameters, cross-referencing rather than repeating the streaming/windowing knobs already listed under Training data order and the sparse all-to-all precision knobs under Precision. Add the LR schedule knobs LR_WARMUP_STEPS and LR_WARMUP_START_LR, which were absent from the README even though section 5 pins their values, and note that DENSE_LR/SPARSE_LR are the base LR the warmup ramps toward - the parameter section 5 calls the target (post-warmup) learning rate. Co-authored-by: Cursor <cursoragent@cursor.com>
The fit through the RCP endpoints reproduced the measured earliest AUC 0.75 crossing exactly at GBS 8192 and 32768, so eval started on precisely that boundary with no headroom. That minimum is a statistic over 20 seeds and can only move down as more seeds are collected: a faster seed would cross the target before the first eval fires and train on to the next boundary, inflating its reported time to converge. Back the start point off by 2 eval blocks: start_samples = FLOOR((4480*GBS + 149094400) / 3) - 2*2293760 giving 0.025 / 0.0303 / 0.041 at GBS 8192 / 16384 / 32768 (was 0.027 / 0.0323 / 0.043), each landing exactly 2 boundaries earlier. Measured cost across the 60 RCP logs is 5.5-6.9% eval share, up from 4.4-5.1%, or under 3 minutes of wall clock - still far below the ~29% with no skipping, and within the ~5-7% the gin comment already documents. The margin is defined in eval blocks (eval_interval_steps*GBS = 2,293,760) rather than the nominal EVAL_EVERY_DATA_PCT*total_train_samples (2,290,835) because the eval grid is rounded to whole steps; the nominal value leaves the margin a fraction of a step short and the CEIL in the threshold resolution rounds that back up into a whole missed eval. Documented in both files. No code change - this only moves the configured values and the derivation. Co-authored-by: Cursor <cursoragent@cursor.com>
The start point was documented as the RCP fit with the safety margin trailing it as a separate term: start_samples = FLOOR((4480*GBS + 149094400) / 3) - 2*2293760 The margin (2 eval blocks = 4,587,520 samples) is an exact multiple of the divisor, so subtracting 3*4,587,520 = 13,762,560 inside the floor is identical to subtracting 4,587,520 outside it. Folding it into the intercept collapses the whole thing to one expression: start_eval_samples = FLOOR((4480*GBS + 135331840) / 3) This is an algebraic identity, not an approximation. Every value is unchanged: 57,344,000 / 69,577,386 / 94,044,160 samples at GBS 8192 / 16384 / 32768, still rounding to the configured 0.025 / 0.0303 / 0.041. Verified against the old form at all three batch sizes on a compute node (job 35363). Also add the RCP sweep figure to README section 5. It plots the min/max/median/ mean epoch of the first AUC 0.75 crossing against the start-eval line, making visible why the start point has to scale with GBS (all four statistics rise), why the fit is built on the min curve, and that the line sits below min at every batch size. Track skip_eval_epoch_pct_test.py, which validates the documented formula against the gin defaults but was never committed alongside the other tests in that directory. Its formula assertion now uses the folded constant; the test that checks the margin still moves the start exactly 2 eval boundaries keeps the pre-margin intercept 149094400, since it needs the bare fit to compare against. No behavior change - documentation and a test, no runtime code touched. Co-authored-by: Cursor <cursoragent@cursor.com>
The MLPerf compliance checker only evaluates a rule when a line carrying that key is present, so a hyperparameter the closed division fixes but the run never logs cannot be enforced at all. Emit the dense Adam coefficients and the LR warmup schedule from log_run_start, reading the same gin bindings the optimizer uses so the log cannot drift from the configuration. Rename the benchmark short name to dlrmv4 to match the string being registered in mlcommons/logging for Training 6.1. Default AUC_THRESHOLD to the 0.75 quality target so a run early-stops on convergence and emits RUN_STOP with status=success. Both checkers need that shape: compliance requires exactly one run_stop, and the RCP checker treats a run without a successful one as never converged. The gin default alone was not enough since launch_slurm.sh passed an explicit 1.0 fallback that overrode it, and run_and_time.sh still carried DLRM DCNv2's 0.80275 target. Co-authored-by: Cursor <cursoragent@cursor.com>
Completes the set of fixed hyperparameters in the MLLog stream. The closed rules can only enforce a value that appears as an event, so anything the README calls fixed but never logs is documented rather than checked. The eval grid matters beyond bookkeeping: samples-to-converge is quantized to eval_every_data_pct blocks, so a submission is only comparable against the RCPs when its cadence matches. skip_eval_epoch_pct is batch-size dependent and easy to leave at the gin default, which is the GBS 8192 value; logging it lets the rules recompute the expected fraction from global_batch_size and reject a mismatch instead of silently shifting where measurement starts. Values are read from the same gin bindings the model resolves, so the log cannot drift from the config. Co-authored-by: Cursor <cursoragent@cursor.com>
The benchmark registers upstream as dlrmv4 for Training 6.1 and the log
stream already emits that name, so the reference implementation should not
still be shaped like v3. Renames the package directory, every import and
documented path, and the DLRMv3* dataset classes.
The class rename is not cosmetic: get_dataset("yambda-5b") returns
DLRMv3YambdaDataset, so the stale name sat directly on the live training
path rather than in unused vendored code. Its base class is shared with the
debug and synthetic datasets, so the whole hierarchy moves together.
rcp_logs/ is deliberately left alone. Its 62580 matches are the "file"
metadata of the 60 frozen sweep runs, recording where the code actually was
when they executed; rewriting them would misstate the provenance of the RCP
numbers extracted from them.
Co-authored-by: Cursor <cursoragent@cursor.com>
Replaces the TODO_GENERATE_HASH placeholders, which made verify_dataset.sh silently degrade to an existence check and warn that checksums were not pinned. Hashes come from the canonical preprocessing output on the shared volume, 47.6 GB across the five files, and verify_dataset.sh now takes the md5sum -c path and passes. The header records the source HuggingFace snapshot and notes that these hash the artifacts byte for byte, so they verify a copy or a download rather than an independent re-run of preprocessing, whose parquet writer would emit different bytes for logically identical data. Co-authored-by: Cursor <cursoragent@cursor.com>
DLRMv4YambdaDataset._load_metadata derives the embedding-table size from the maximum item_id across artist_item_mapping.parquet and album_item_mapping.parquet, so a truncated or stale mapping does not raise. It silently changes the model being benchmarked, which is exactly the failure a checksum should catch. The two files add 93 MB to a 47.6 GB verification. They are addressed relative to the processed dir so a single md5sum -c run still covers everything. Co-authored-by: Cursor <cursoragent@cursor.com>
…length _load_metadata looked for item_popularity.npy in metadata_dir, but preprocessing writes it next to the sessions in processed_dir. The lookup always missed and fell back to deriving the vocabulary size from the artist/album mappings alone. On the current dataset all three sources agree at 9,390,624 items, so the fix is bit-wise equivalent: the derived n_items and both item_to_artist / item_to_album arrays are byte-identical before and after, which keeps the frozen RCPs valid. Reading the file from the right place also makes it useful as a cross-check, so a length that disagrees with the mappings now raises instead of silently resizing the embedding tables. Co-authored-by: Cursor <cursoragent@cursor.com>
The dataset can now be obtained two ways: a prepared download that ships the derived flat-event cache and skips preprocessing, or the existing build from the upstream HuggingFace release. They deliver different files -- the prepared copy carries no parquet, since train_sessions.parquet is only the input the cache is built from -- so verify_dataset.sh no longer assumes one fixed file set. It now requires the two item mappings plus at least one trainable source, and checks whatever else is present. Cache checksums are sha256 in a separate file because they verify a locally built cache as well as a downloaded one. The build is deterministic: sessions sort by (uid, session_id), a total order, and .npy stores only dtype and shape. Rebuilding from the same train_sessions.parquet on a different host reproduced all sixteen files byte-for-byte. Co-authored-by: Cursor <cursoragent@cursor.com>
The dataset is now published on the MLCommons training data page, so the placeholder rclone-with-credentials recipe is replaced by the standard R2 downloader. It resolves the file list from dlrmv4_dataset.uri and md5-verifies each file as it downloads, which removes the credential handling and the separate manifest step. Co-authored-by: Cursor <cursoragent@cursor.com>
The header tells readers they can run md5sum -c directly, which holds for a local build but reports four failures on a prepared download: that layout ships only item_popularity.npy and the two mappings of the seven files listed. Point those readers at verify_dataset.sh, which checks what is present. Co-authored-by: Cursor <cursoragent@cursor.com>
The FLOPs estimate is the MFU denominator, but nothing recorded why it charges x3.5 and a causal x1/2 on attention or x3 on the weight GEMMs, so the numbers were hard to audit and easy to misread as disagreeing with the published DLRMv3 inference formula. Spell out the per-component derivation next to the code, including which weight shapes confirm the 4x UVQK and 3x output-projection widths, and add the same breakdown to the README so reported TFLOPs can be reconciled term by term. Multipliers now come from named constants rather than bare 6 and 7. Co-authored-by: Cursor <cursoragent@cursor.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
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a new DLRMv4 MLPerf Training reference benchmark: an HSTU
(Hierarchical Sequential Transduction Units) generative recommender trained on
the public Yambda-5b dataset, vendored under
recommendation_v4/as asibling of the existing
recommendation_v2/torchrec_dlrm.What this adds
generative_recommendersdlrm_v3 path) onTorchRec
DistributedModelParallel, bf16 training, singlelisten_plustask. Architecture,
history_length,hbm_cap, and embedding dims aregin-tunable with env overrides.
preprocess_public_data.pydownloads + preprocessesYambda (50m / 500m / 5b) from HuggingFace (temporal GTS split, session
segmentation, item popularity);
DLRM_DATA_PATHenv override.no future leakage. Eval is a strictly-future held-out window — either the next window (T → T+1)
or a fixed holdout window/dataset (eval_holdout_ts, optional train/eval split).
Persistent-loader + double-buffer + eval-prefetch hide window-reset overhead;
crash-resumable checkpoints with step/time/window cadences and bit-equal RNG replay.
(TMA early-out, jagged multirow routing, separated-RNG LN-dropout, autotune
pinning) plus NVIDIA B200 (sm_100). Multi-node SLURM launcher with GPUDirect
RDMA.
mllogintegration centralized inmlperf_logging_utils.py(resume-aware INIT/RUN/EVAL markers, AUC convergencetarget 0.80275, AMD/MI355X submission identity);
mlperf_loggingpinned to6.0.0-rc6 in Dockerfiles + requirements.
download_dataset.sh/verify_dataset.sh/run_and_time.shwrappers, MLPerf-spec README (summaries, model+paper,hyperparameter table, quality target, eval frequency), frozen
requirements.txt, and an RCP placeholder.dependency versions, perf ladder) and perf/multi-node notes.