Testing - #902
Closed
ShriyaRishab wants to merge 7 commits into
Closed
Conversation
- Introduced a default PR template for new benchmarks, outlining necessary steps and checklist items. - Added a separate PR template for introducing new rounds, including confirmation of round scope and updates to relevant documentation. - Updated CONTRIBUTING.md to reference the new templates and provide guidance on their usage. Signed-off-by: ShriyaRishab <spalsamudram@nvidia.com>
* Add recommendation_v4 (HSTU/DLRM-v3 generative-recommenders fork @ d97e51c)
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>
* Enable Triton HSTU kernels on AMD/ROCm (gfx950 MI350X)
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>
* Fix AttributeError on triton.knobs.nvidia.use_meta_ws
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.
* Make HSTU model arch + dataset history_length gin-tunable
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 EmbeddingShardingPlanner hbm_cap_gb gin-tunable, set ddr_cap=0
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>
* Add yambda 50m/500m/5b preprocessor + DLRM_DATA_PATH env override
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>
* Fix NCCL init: set CUDA device before init_process_group
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.
* Profiler: write traces to local disk under <repo>/results/<run_name>/
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>
* Tracing fixes: gin scoping, drop active=10 override, intuitive filenames, 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>
* gin: history_length 2048 → 2039 + expanded per-pool comment
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>
* README: rewrite for yambda-5b fork — upstream link, data prep, per-pool 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>
* gin: make hbm_cap_gb overridable via \$HBM_CAP_GB
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>
* docs: add B200 training recipe for yambda-5b
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.
* bf16 + triton autotune pinning with gin-driven full-tune override
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>
* docs: add MI350X training recipe section
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>
* scripts: add stitch_traces.py
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>
* docs: update B200 recipe deps to NGC 26.04 (torch 2.12 / CUDA 13.2)
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.
* docs: refresh B200 recipe deps (fbgemm HEAD, torchrec 1.7 nightly, driver)
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).
* MI350X: re-pin 2 triton configs for the torch 2.12 + torchrec 1.7 stack
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>
* docs: update MI350X Stack B to fbgemm @ B200 commit + caveat
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>
* docs: drop Stack A; MI350X recipe is now single-stack (B200-aligned)
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>
* docs: drop PYTORCH-fallback caveat from MI350X recipe
Not relevant — TRITON is the documented default backend.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* MI350X: fit-entity embedding sizes, bs=1024 default, batch-agnostic recipe
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>
* MI350X: separated-RNG LN-dropout + attention autotune pin + clock guard
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>
* dlrmv4: TorchRec 3-stage sparse-dist pipeline + gin-selectable HSTU kernel
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>
* dlrmv4: streaming (temporal-order) training for yambda-5b
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>
* dlrmv4: disable checkpointing by default; fix recipe torch note
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>
* dlrmv4: ROCm-only Perfetto trace render fixes at export time
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>
* dlrmv4: ROCm annotation de-overlap so phase spans render full width
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>
* dlrmv4: streaming checkpoint resume + step/time checkpoint cadences
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>
* dlrmv4: move streaming resume test harness into train/tests
It's a test driver, not a general script — colocate the shell harness with
its Python comparator under train/tests/ and fix the stale path references.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: sparse full-holdout eval cadence + eval-pool fork-race fix
Add eval_every_n_windows (env EVAL_EVERY_N_WINDOWS, default 1 = no change) so
the heavy full next-day eval window can run every Nth train window (and always
the final one) instead of every window, amortizing its cost on the long run.
Fix a deadlock this exposed in the double-buffer path: the persistent eval
worker pool's first iter() (its only fork) must happen on the main thread
BEFORE the prefetcher's background prep thread starts. Deferring that first
fork into the loop (as the sparse cadence naively did) forks while the bg
thread holds an allocator/GIL-released lock and hangs the run. Always pre-fork
the eval pool before the loop; in-loop re-arms only reset the persistent
workers (no fork) and target the next window that will actually eval.
Also normalize NUM_TRAIN_BATCHES/NUM_EVAL_BATCHES <=0 to None (full window /
full-holdout eval) and bind NUM_EVAL_BATCHES in gin so eval can be capped for
fast validation without affecting the full run.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: self-healing streaming-e2e supervisor + NE/AUC trajectory builder
run_streaming_e2e.sh: head-node supervisor that keeps a multi-day yambda-5b
streaming train+eval alive across (1) trainer crash/OOM, (2) silent SIGKILL,
and (3) node loss. Relaunches from the latest checkpoint each time (exact-once
resume handles continuity). Node failover salloc's a fresh exclusive node on
the partition, provisions the container on it, and resumes from shared NFS;
allocations it creates are released on success (never the user's own --jobid).
Includes disk guard + stale .tmp sweep, keep_last_n retention, an exit-sentinel
+ stall watchdog for crash detection, and a node-health watchdog. Heavily
documented inline.
build_ne_auc_trajectory.py: parse train+eval NE/AUC (+perf) from a run log and
emit combined CSV/JSON plus an NE/AUC-vs-step trajectory plot.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: durable streaming-run metrics (append log + TensorBoard on NFS)
Make the full-run NE/AUC record survive relaunches and node failover:
- launch_smoke_8gpu.sh now appends to $LOG (tee -a) instead of truncating, so
a supervised run that relaunches many times into the same log keeps its full
metrics history. The supervisor initializes the log once at run start.
- run_streaming_e2e.sh: truncate $LOG once at start, create the per-run TB dir,
and export TENSORBOARD_LOG_PATH=/apps/chcai/tb/$RUN_NAME/ into the launch env.
- yambda_5b.gin: MetricsLogger.tensorboard_log_path now reads $TENSORBOARD_LOG_PATH
(via the existing env_path helper) defaulting to /apps/chcai/tb/yambda_5b/ on
shared NFS, instead of container-local /tmp (which is wiped on failover).
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: raise streaming-run disk guard for keep_last_n=1 saves
A checkpoint save writes a fresh ~560 GB .tmp before the old copy is
pruned, so peak transient usage is (keep_last_n + 1) copies (~1120 GB at
keep_last_n=1). Bump MIN_FREE_GIB 800 -> 1200 so a (re)launch never wedges
mid-save on a near-full shared NFS volume.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: anchor eval points to train global step in NE/AUC trajectory
Eval used a per-window-resetting internal step counter, so eval markers
bunched at the left of the x-axis instead of overlaying the train curve.
Parse the log sequentially, collapse each eval window to its final
full-holdout metrics, and anchor it to the train global step it ran at
(tagging eval_ts via the [boundary] marker, which can interleave around
the eval's metric lines). Eval points now overlay train on a shared axis.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: supervisor tolerates control-plane outages + attach mode
A SLURM controller outage made the supervisor's squeue/sinfo health check
report "node lost", trigger failover, and then FATAL-exit when salloc
couldn't reach the controller either - abandoning a run whose trainer was
in fact still alive (slurmctld outages don't kill RUNNING jobs).
Harden it:
- controller_up()/wait_for_controller(): treat an unreachable controller as
transient; wait for recovery (up to --ctrl-wait-max) instead of failing over.
- Direct-SSH fallback (dexec via cached LAST_NODE) so trainer_alive() and the
mid-run watchdog verify liveness even while the controller is down; only fail
over if the trainer is genuinely gone, not on a control-plane blip.
- timeout-guard all srun/squeue/sinfo calls so a hung control plane / NFS can't
wedge the supervisor.
- --attach mode: adopt an already-running trainer (one that outlived a killed
supervisor) without truncating its log, sweeping its in-flight .tmp, killing
it, or relaunching - just resume monitoring in place.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: TFLOPS/MFU/HFU reporting + MAX_SEQ_LEN/HISTORY_LENGTH gin knobs
Adds per-step TFLOPS, MFU and HFU to the training perf line + TensorBoard
scalars, plus env-driven gin knobs for sequence-length sweeps. The two
changes ship together because TFLOPS reporting is what makes the 4k vs
2k comparison interpretable (a sps drop alone doesn't say whether the
GPU is doing more work per sample or less work overall).
TFLOPS reporting
HSTU is structurally different from a standard transformer: the UVQK
projection fuses Q/K/V/U into one matmul, and SiLU(U)*y elementwise
gating replaces the FFN block. So a TorchTitan-style "count matmuls
with factor 6, attention with factor 12" template gives the wrong
per-sample FLOPs unless it's coded against HSTU's specific shapes.
DlrmHSTU.get_num_flops_per_sample() implements the HSTU-specific
dense formula (UVQK projection + Q.K^T + att.V + output projection,
per layer, times n_layers). Multitask head adds a trivial constant.
Embedding lookups excluded because they're memory-bound and would
otherwise pollute MFU. This is the dense yardstick: what the FLOPs
would be if every sample's UIH filled max_seq_len. It's the standard
MFU denominator (matches Primus-DLRM's OneTrans accounting style).
Yambda's per-user history is jagged, so the actual GPU work is
significantly less than the dense estimate. main_forward stashes
_last_jagged_flops_per_sample after computing each batch's mean(s)
and mean(s^2), and MetricsLogger reads + .item()s it once per
metric_log_frequency (one D->H sync per logging interval, not per
step). When present, the perf line splits into:
tflops_algo/gpu mfu - dense yardstick, MFU denominator
tflops_real/gpu hfu - actual jagged work, hardware utilization
fill - real / algo, padding-skipped fraction
When the jagged stash is absent (other model types, or before the
first main_forward), only tflops_algo/mfu print. When the model
doesn't expose get_num_flops_per_sample at all, the perf line is
byte-for-byte unchanged (backward compatible).
get_gpu_peak_flops("bf16"/"fp32") consults a per-GPU peak table
(MI355X/MI350X=2300 TF, MI300X/MI325X=1300, B200=2250, H100=990,
A100=312 for bf16) and warns + defaults to MI350X for unknown
device names. train_ranker pulls "bf16" when bf16_training=True
else "fp32"; the dtype string drives only the denominator, not
anything else.
TensorBoard scalars added alongside the existing perf/* group:
perf/train_tflops_algo_gpu, perf/train_mfu_pct,
perf/train_tflops_real_gpu, perf/train_hfu_pct,
perf/train_fill_pct.
Validated on 8x MI350X yambda-5b at the 2k baseline:
241.6 GFLOP/sample (dense), GPU peak 2300 TFLOPS
steady-state ~13500 sps -> mfu 17.7-17.9%, hfu 9.7-10.2%,
fill 55-57% (yambda users average ~1170 events vs 2046 max).
HSTU's MFU is higher than the OneTrans baseline on the same GPU
(Primus-DLRM hit 5.7%) because HSTU does less compute per token,
so the FLOPs it does execute run at higher hardware utilization.
MAX_SEQ_LEN / HISTORY_LENGTH env knobs
Adds two env-driven gin macros so sequence-length sweeps don't
require editing yambda_5b.gin (which a running e2e job has parsed
-- editing the file would change behavior on a supervisor restart):
get_hstu_configs.max_seq_len = @msl/env_int() default 2048
get_dataset.history_length = @hl/env_int() default 2039
Defaults are the current production values, so unset env is a no-op.
Used the 4k validation run with MAX_SEQ_LEN=4096 HISTORY_LENGTH=4096
(reuses hstu_cache_L4096/ on disk; ~8 events of trailing UIH
truncation per sample, negligible).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* dlrmv4: configurable lifetime-AUC backend + fixed-holdout streaming eval
Add a dual-set streaming eval (fresh per-pass "window_*" + cumulative
"lifetime_*") for NE/Accuracy/GAUC/AUC, with a gin-selectable lifetime-AUC
backend for both train and eval: "binned" (BinnedCumulativeAUC, exact
cumulative AUC from an O(bins) histogram, default) or "capped"
(LifetimeAUCMetricComputation, trailing per-rank buffer). Backend, bins, and
window are MetricsLogger gin bindings (env-overridable).
Persist per-rank cumulative metric state in metricbuf_rank{rank}.pt for both
backends across train/eval/eval_cum; keep per-rank state out of the shared
rank-0 blob (strip capped buffers, zero binned histograms) so a resume never
inherits rank-0's counts. Eval set is a stable user-hash holdout over a fixed
window range, validated against a checkpoint split contract on resume.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: default yambda-5b to the 4k-no-truncation seq shape
Set HISTORY_LENGTH=4086 / MAX_SEQ_LEN=4096 as the gin defaults (3*1362+9=4095
≤ 4096, the no-overfill 4k analog of the prior 2039/2048 shape). Override via
$HISTORY_LENGTH/$MAX_SEQ_LEN; use 2039/2048 to reuse the 2k single-task cache.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: reservation-aware sbatch failover in streaming supervisor
Replace the salloc-based node failover with an sbatch hold job (--wrap "sleep
infinity", bounded by --time), since interactive salloc on meta64 is capped at
240 min and can't hold a multi-day run. Add --reservation so a replacement node
is re-acquired from the same SLURM reservation, plus --acquire-wait-max to
tolerate brief queueing. Provisioning still runs via srun --overlap afterward.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: two-tier reservation-then-open-pool failover in streaming supervisor
acquire_node now prefers the configured --reservation (tier 1, short
RESV_WAIT_MAX wait since a free reservation node starts ~immediately),
then falls back to the open partition pool (tier 2, ACQUIRE_WAIT_MAX).
The pending reservation hold job is scancel'd before the fallback resubmit
so we never end up holding two nodes. Factors the sbatch submit and
RUNNING-wait into _submit_hold_job/_wait_running helpers; adds the
--resv-wait-max knob.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: cap failover at <=1 reservation node + fix trainer-alive self-match
trainer_alive: `pgrep -f generative_recommenders` always matched the probe
shell's own cmdline, so it could never report the trainer dead -- defeating
the stall watchdog and making ATTACH mode falsely "adopt" a nonexistent
trainer. Use the `set -f; pgrep -f [g]enerative_recommenders` self-match guard.
Reservation cap: every node we acquire is an sbatch --job-name=e2e_failover
hold, so reap_failover_holds() reaps strays by name at startup (catching holds
leaked by a prior supervisor that died mid-failover) and before every acquire
(no stacking). wait_for_original_recover() waits --orig-recover-wait (def 600s)
for SLURM to requeue the lost ORIGINAL job and reuses it instead of grabbing a
SECOND reservation node. Together these keep us at <=1 reservation node.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: anchor sparse-eval cadence to absolute window ts (resume-invariant)
_should_eval keyed the every-N-windows cadence off the per-call loop index
`i`, so a mid-run resume (which rebases start_ts and restarts train_ts_list at
the resume window) re-anchored the eval grid -- e.g. evals shifted from
150,160,170,... to 165,175,185,... after resuming at window 165. Capture the
original start_ts as eval_anchor_ts BEFORE the resume block mutates start_ts,
and gate eval on (train_ts_list[i] - eval_anchor_ts) % K == 0, so the eval grid
is identical on cold start and every resume. Final-window eval preserved; the
eval-pool fork is unconditional so dropping the i==0 eval on resume is safe.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: multi-node (N>=1) training over RoCE RDMA via consolidated launch_slurm.sh
Consolidate the SLURM orchestration, container/RDMA provisioning, and
in-container trainer launch into a single self-dispatching scripts/launch_slurm.sh
(phases: orchestrate -> provision -> worker) supporting N>=1 nodes; N=1 keeps the
legacy single-node path byte-for-byte.
Multi-node runs real RDMA over the 8 Broadcom bnxt_re RoCE HCAs. The key fix is
an LD_PRELOAD/LD_LIBRARY_PATH overlay of the host's matched rdma-core (v61/v59):
the container's stock v34 provider faults RCCL's deep create_qp (256 WRs) against
the host kernel uapi -> "ibv_create_qp ... Bad address". The unversioned
libibverbs.so symlink in the overlay is required so torch maps only the host lib.
TCP bootstrap is pinned to the routable fenic0; RDMA data goes over bnxt_re
(GID idx 3, TC 104).
Python: derive the global rank (node_rank*gpus_per_node+local_rank), forward
master_addr, pass local_world_size to the TorchRec planner and live world_size to
metrics. Single-node behavior unchanged.
Docs: add docs/multi_node_config.md (enablement details, lessons, and the
cluster-specific knobs to change per fabric); README + perf_opt updated for the
launch_slurm.sh rename.
Co-authored-by: Cursor <cursoragent@cursor.com>
* local dlrmv4 changes: docker setup, run scripts, walkthrough docs, smoke log path
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: decouple multi-node launch from per-user paths for portable baseline
Run the consolidated launch/streaming flow from any $HOME without editing
another user's tree:
- Derive REPO_MOUNT/DATA_MOUNT/SCRATCH from $HOME with env overrides; keep
the shared read-only dataset path intact (no data duplication).
- Per-user container name (yambda_$USER) to avoid collisions.
- In-repo provisioning via launch_slurm.sh (drop external _provision script dep).
- chmod log files world-writable so the container (nobody) can tee-append
under NFS root_squash; fixes spurious pipefail rc=1.
- Neutral in-repo TensorBoard default path.
Validated: 2-node ROCm/RDMA smoke (world_size 16) completed rc=0, 20-batch
train window with metrics logged.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: deterministic in-window shuffle diversity dial + opt-in diagnostics
Add STREAMING_SHUFFLE_FRACTION (0..1) as a config-invariant control over
in-window embedding diversity, replacing the legacy block/buffer knobs. Full
shuffle is the deterministic, seeded default (fraction=1.0). Add an opt-in
unique-embedding diagnostic (DIAG_UNIQUE_EMB) and gate chrome-trace capture
behind OUTPUT_TRACE; both default off to keep production runs overhead-free.
Forward the new env knobs through launch_slurm.sh.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: MLPerf training compliance logging for streaming-train-eval
Wire mlperf_logging (mllog) into the yambda-5b streaming-train-eval path:
rank-0-gated MLPerfLogger facade emitting the full event lifecycle
(cache_clear/init_start -> submission_info + hyperparameters ->
init_stop/run_start -> per-window block_start/stop + eval_start/accuracy/stop
-> run_stop), driven by a configurable convergence target.
Key points:
- AUC_THRESHOLD (gin, env-overridable; default 0.80275) doubles as the MLPerf
convergence target: rank 0 decides on the global lifetime eval AUC and
BROADCASTS the stop boolean so all ranks break in lockstep (avoids the
ALLTOALL collective-timeout deadlock from a per-rank decision).
- MLPerfLogger uses the explicit global rank passed by train_ranker (computed
pre-dist-init), so only true rank 0 logs and the compliance file has exactly
one event each. Per-job MLPERF_LOG_PATH avoids stale append accumulation.
- Per-step train_loss POINT_IN_TIME event (global cross-rank mean) with
samples_count + lr, plus a console/TensorBoard readout.
- cumulative_train_samples counter (global, checkpointed) as the samples_count
progress unit; lifecycle gated on cold start so e2e-supervisor resumes never
emit orphaned run boundaries.
Validated end-to-end at 1/2/4 nodes (8/16/32 GPUs): clean compliance log,
single run_stop=success, passes the common/closed_common compliance checks.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: min_history anchor-eligibility floor (decoupled from history_length)
A LISTEN event qualifies as a train/eval anchor once the user has >= min_history prior events, decoupled from history_length (the gather/truncation cap) since jagged attention handles short UIH. Anchor positions/anchor_ts caches are keyed by (history_length, min_history) and built independently of the _READY-gated 150GB flat store, so changing the floor rebuilds only the cheap positions array. Default None preserves the legacy full-history behavior (which dropped ~60% of users).
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: gin defaults — min_history=1, BATCH_SIZE env knob, default to no in-window shuffle
Bind get_dataset.min_history to $MIN_HISTORY (default 1 = ~all users). Make batch_size env-overridable via $BATCH_SIZE (default 1024). Change default streaming_shuffle_fraction 1.0 -> 0.0 (user-major, production streaming order); override per-run via $STREAMING_SHUFFLE_FRACTION (1.0 = full shuffle).
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: launch_slurm container hygiene + readiness gating + env passthroughs
Reap stale foreign GPU containers and restart our own to reclaim leaked HBM before launch; gate the worker exec on container State.Running + a probe with retry to fix container-restart races (which caused NCCL TCPStore 600s timeouts); APPEND_LOG=1 appends the metrics log on resume. Forward MIN_HISTORY/MAX_SEQ_LEN/HISTORY_LENGTH/BATCH_SIZE/CKPT_TIME_INTERVAL_S/DIAG_EMB_STEPS into the container. Drop the hardcoded --time; tidy header comments after launch_smoke_8gpu removal.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: durable per-boundary eval metrics JSONL sink + aggregate emb-diag
Write one JSON line per eval boundary to <LOG>.metrics.jsonl capturing the end-of-pass metrics over the fixed holdout (append-only, rank 0, survives restarts/resumes) — no interim/averaging ambiguity. Rework the unique-embedding diagnostic into an aggregate over DIAG_EMB_STEPS batches covering the cross-feature tables.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: default MIN_HISTORY=0 (include cold-start first events) to match min0 runs
All live min0 runs already export MIN_HISTORY=0; make the gin default match so future runs without an override include each user's cold-start first event (zero prior context) instead of requiring >=1 prior event.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: default NUM_TRAIN_TS=149 to sweep full ts=150..298 streaming range
Matches the long e2e runs (start_ts=150 + 149 daily windows). EVAL_EACH_WINDOW
and EVAL_EVERY_N_WINDOWS already default to 1. Clamped to the dataset's
available window count at runtime; override via $NUM_TRAIN_TS.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: consolidate eval cadence into single EVAL_EVERY_N_WINDOWS knob
Drop the redundant EVAL_EACH_WINDOW on/off boolean and fold "disable eval"
into the cadence int: 0 = eval off (train-only / resume test), 1 = every
window, N>1 = every Nth window (anchored to the absolute ts grid, stable
across resume). Updates streaming_train_eval_loop, the gin bindings, and the
resume test harness accordingly. No behavior change for the live runs
(EVAL_EVERY_N_WINDOWS=1).
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: random per-run seed + graceful teardown + no double-logging
- setup(): init the process group first, then draw a fresh random seed per
run (rank 0 broadcasts so all ranks agree), export it to $SEED, and log it;
pin $SEED to reproduce a run exactly. Data order/split are unaffected (still
time-deterministic + $SPLIT_SALT-governed); the seed governs dense weight
init. MLPerf SEED event now logs the actual chosen value, not hardcoded 1.
- train_ranker(): move distributed teardown into a finally block so a clean
finish also barriers + destroy_process_group()s in lockstep, silencing the
noisy TCPStore "broken pipe" / "should dump" warnings at exit. Best-effort
so teardown never masks a real error.
- launch_slurm.sh: orchestrate sets WORKER_TEE=0 so the worker points its own
file sink at /dev/null (stdout is already tee'd upstream), avoiding every log
line being written twice. Direct single-node (e2e supervisor) keeps WORKER_TEE
unset and still writes $LOG itself.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: fix single-node-without-slurm via run_docker.sh
run_docker.sh: strip an optional leading `--` so the documented
`run_docker.sh -- bash scripts/launch_slurm.sh` form works (it was
forwarded verbatim to `docker run` and failed exec: "--" not found);
wire run_docker.sh to the launch_slurm worker flow (CONTAINER_NAME,
LOG/MODE/MAX_SEQ_LEN/HISTORY_LENGTH passthrough) and forward
NCCL_SOCKET_IFNAME so the bootstrap NIC is host-overridable.
launch_slurm.sh: default NCCL_SOCKET_IFNAME=lo for single-node
(NNODES==1) instead of the meta64-only fenic0 — loopback is reachable
by all local ranks on any host (data plane is intra-node XGMI/PCIe),
so the single-node path now runs out-of-the-box on dev boxes with no
fenic0. Multi-node keeps the fenic0 default; both stay overridable.
Verified: 8-GPU single-node streaming-train-eval smoke runs clean
(RCCL init over lo, train+eval, MLPerf run_stop) on a lone MI355 node.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: gin-configurable quantized (bf16/fp16) embedding all-to-all
Quantize the bandwidth-bound embedding-shuffle all-to-all via TorchRec
QCommsConfig on the sequence EmbeddingCollectionSharder. Exposed as two
gin knobs on make_optimizer_and_shard (env-overridable):
sparse_a2a_precision = fp32 (off, default) | bf16 | fp16
sparse_a2a_quantize_backward = 1 (default) | 0 (forward-only)
Default fp32 keeps the path byte-for-byte identical to trunk. launch_slurm
forwards the env overrides and reuses a stopped container instead of
destructively re-provisioning.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: split quantized a2a into independent fwd/bwd precision knobs
Replace the single sparse-a2a precision knob with separate forward and
backward precision settings ($SPARSE_A2A_FWD / $SPARSE_A2A_BWD, each
fp32|bf16|fp16; both fp32 = off, identical to baseline). This enables the
TorchRec golden_training recommended mix (fwd=fp16, bwd=bf16): fp16's
mantissa suits bounded forward activations while bf16's wider exponent
range avoids gradient overflow. 2-node A/B shows fp16/bf16 at perf parity
with bf16/bf16 (both 2-byte wire), so it's a numerical-safety win at zero
perf cost.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: enable GPUDirect RDMA by default in slurm worker
Set NCCL_NET_GDR_LEVEL=5 and NCCL_DMABUF_ENABLE=1 by default so RCCL does
true GPU<->NIC DMA over bnxt_re instead of host-memory staging. The brcmrdma
host kernel ships the inbox peer-memory client, so GDR works with no
container/host changes. Measured ~+22% throughput at 2 nodes (65.7%->79.8%
weak-scaling efficiency). Overridable via NCCL_NET_GDR_LEVEL=0 and non-fatal
(falls back to host staging if peermem is absent).
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: gin-configurable RNG seed ($SEED) + default TRAIN_SPLIT_PERCENTAGE=1.0
Add a gin-configurable global seed so weight init (dense params + embedding
tables) is reproducible run-to-run and runs can be init-matched A/Bs. The seed
is bound via $SEED (seed_everything.seed = @seed/env_int(), default 1) and
applied by a new seed_everything() called in train_ranker right before
make_model() — after the full gin parse, so the binding resolves in the second
parse where env_int is registered. Move the old hardcoded seed=1 out of setup()
(too early to be gin-configurable). Forward $SEED through launch_slurm.sh.
Flip the default TRAIN_SPLIT_PERCENTAGE 0.90 -> 1.0 (all users trained AND
evaluated, matching the alleval/qa2a production runs) in both the gin default
and the launch_slurm.sh fallback.
Validated with two short 1-node runs: SEED=1 and SEED=2 each log their seed on
all ranks, and tsp=1.0 is applied without exporting TRAIN_SPLIT_PERCENTAGE.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: env-configurable dense/sparse LR + optimizer LR logging
Make the dense ($DENSE_LR) and sparse ($SPARSE_LR) optimizer learning rates
overridable per-run via env (defaults unchanged at 0.001), and log the
resolved LR at optimizer construction so runs are self-documenting. Forward
both vars through launch_slurm.sh into the container.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: env-configurable HSTU transformer depth ($HSTU_NUM_LAYERS)
Make the HSTU attention layer count overridable per-run via $HSTU_NUM_LAYERS
(default 5, unchanged), resolved in the full gin parse and forwarded through
launch_slurm.sh. Changing depth alters model shape, so a run with a new depth
must use a fresh CKPT_PATH (incompatible with existing checkpoints).
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: env-overridable dense/sparse LRs for sweeps + holdout default 1.0
Make dense (Adam) and sparse (RowWiseAdagrad) learning rates overridable via
$DENSE_LR / $SPARSE_LR with gin defaults preserved at 0.001, so LR sweeps don't
require editing gin. Resolve gin macro references in the MLPerf param logger so
env-overridden LRs are logged as real numbers. Default TRAIN_SPLIT_PERCENTAGE
to 1.0 (no holdout) and log the resolved LR overrides at orchestration time.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: disable TensorBoard by default (no-op writer)
The shared-NFS tfevents writer was the only metrics sink whose I/O error
was uncaught, and it repeatedly crashed trainers on transient /apps
Errno 121 (Remote I/O) hiccups. Default TENSORBOARD_LOG_PATH is now empty,
which installs a _NoOpSummaryWriter so the metrics path (compute +
text-log + .metrics.jsonl sinks) runs unchanged and never crashes on TB
file I/O. Nothing we consume reads TensorBoard. Re-enable per-run by
setting $TENSORBOARD_LOG_PATH to a non-empty dir.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: default dense/sparse LR=1e-5 and HSTU depth=3
Change the gin defaults to the configuration validated by the recent
power-user (min4086) ht299 runs: dense+sparse LR 0.001 -> 1e-5 and HSTU
attention depth 5 -> 3 (reaches ~0.78-0.81 holdout-299 AUC at ~1.4x
faster training than the 5-layer/5e-5 setup). Both remain env-overridable
via $DENSE_LR / $SPARSE_LR / $HSTU_NUM_LAYERS.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: report EVAL_ACCURACY as per-window AUC (configurable, default window)
The MLPerf EVAL_ACCURACY event and the convergence decision (early SUCCESS
RUN_STOP + end-of-run finalize) now use the per-pass full-holdout "window_auc"
instead of the cumulative "lifetime_auc". Made selectable via a new gin knob
streaming_train_eval_loop.eval_accuracy_auc_mode ($EVAL_ACCURACY_AUC_MODE),
default "window"; set "lifetime" to restore prior behavior. Both AUCs are still
computed and logged to TensorBoard. The rank-0-decides-then-broadcast deadlock
guard is preserved. Forward $EVAL_ACCURACY_AUC_MODE through launch_slurm.sh.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: consolidate streaming e2e supervisor to one sbatch-wrapping script
Replace the single-node docker-exec supervisor with the sbatch-job-level
model (formerly run_streaming_e2e_multinode.sh), which handles 1..N nodes
via launch_slurm.sh. Node replacement is now SLURM's job on resubmit, so
the in-place node-acquisition/provision/exec-sentinel logic is dropped.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: add non-SLURM local launcher + self-healing supervisor
launch_local.sh: single-host, scheduler-free analog of launch_slurm.sh's
worker phase — same train_ranker entry point / yambda_5b.gin config, smoke
or full run on a GPU host with no SLURM/docker/RDMA overlay.
run_streaming_e2e_local.sh: local analog of run_streaming_e2e.sh. Backgrounds
a per-run submit script (so wait $PID yields the trainer's real exit code),
relaunches from the latest checkpoint on crash/nonzero-exit/hang, with a
hang watchdog (frozen log + no trainer proc) and a pre-launch disk guard.
cleanup_container waits for GPU HBM to actually drain before relaunching so an
OOM/crash can't cascade into a dirty-GPU OOM loop.
* dlrmv4: set default-PG timeout to TIMEOUT (1800s) to survive checkpoint skew
The checkpoint DCP collectives run on the default process group created by
init_process_group, which had no explicit timeout and thus used NCCL's stock
600s watchdog. The 560GB sparse-embedding checkpoint is written to shared NFS
with a badly imbalanced sharding plan (per-rank shards ~37GB..~95GB), so the
fastest rank can wait >600s in the post-write allgather/barrier for the slowest
rank, tripping the watchdog and SIGABRTing an otherwise-healthy job (observed
on 3 nodes, always rank 7). Pass timeout=timedelta(seconds=TIMEOUT) so the
default PG tolerates the skew, matching the secondary new_group.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: optional gradient clipping for the streaming path ($GRAD_CLIP_NORM)
Add env-configurable global-norm gradient clipping to streaming_train_eval_loop,
applied to dense params after backward() and before optimizer.step() (sparse
tables use a fused optimizer and are unaffected, matching the non-streaming
path's clip_grad_norm_). Wired via gin to $GRAD_CLIP_NORM and forwarded into the
container by launch_slurm.sh. Default 0.0 = OFF, so existing streaming runs are
unchanged. Eliminates the window-131 eval-AUC dip at LR=1e-5.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: data-fraction eval cadence + lr1e-7/grad-clip-on defaults
Add a data-percentage-based eval cadence for the streaming loop, an
alternative to the per-window cadence. EVAL_EVERY_DATA_PCT>0 runs the
full-holdout eval every fixed FRACTION of the run's total training data,
so eval points are evenly spaced by data volume regardless of per-window
sample counts. The fraction is converted once into a global train-step
interval (round(pct * total_train_anchors / (batch_size*world_size)))
over the original requested window range, and eval fires on
global_step % interval -- mid-window and resume-stable, mirroring
checkpoint_step_frequency. Each eval label carries @step=<global_step>
for plotting against data volume.
- yambda: total_train_anchors(start_ts, num_ts) one-time O(N) count.
- streaming_train_eval_loop: eval_every_data_pct param + interval calc;
the two cadences are mutually exclusive (ValueError if both >0).
- launch_slurm.sh: forward EVAL_EVERY_DATA_PCT into the container.
Also flip the yambda-5b gin defaults to the validated config: dense and
sparse LR 1e-5 -> 1e-7, and gradient clipping ON by default
(GRAD_CLIP_NORM 0.0 -> 1.0). HSTU depth stays 3.
Verified e2e on an open-pool node: interval computed correctly, two
mid-window evals fired on the step grid + final eval (rc=0), and the
both-enabled config raises the expected ValueError.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: seed embedding init + reproducibility checksum ($SEED)
Make embedding-table init a deterministic function of $SEED (dense init
already was), so runs are a clean init-matched A/B when data order
($STREAMING_SHUFFLE_SEED) and the holdout split ($SPLIT_SALT) are fixed.
- configs.py: attach a per-table seeded uniform init_fn (per-table seed =
sha256($SEED, table_name)); meta-safe (skips the meta device DMP builds
the unsharded module on). Init bounds mirror stock (+/-1/sqrt(N) or a
table's explicit weight_init_min/max), so the distribution -- and thus
model quality -- is unchanged; only determinism/seeding differs.
- utils.py (make_optimizer_and_shard): re-seed torch/torch.cuda from $SEED
right before DistributedModelParallel(...) so the fused FBGEMM TBE
on-device embedding init is reproducible for a FIXED sharding plan
(Tier 1). Dense params are already built in make_model, so untouched.
- utils.py: log a post-DMP init checksum (per-table count/sum/sumsq +
a one-line digest; sharded stats are all-reduced so the fingerprint
covers the whole table regardless of shard layout). Gate via
INIT_CHECKSUM (default on). Verified on idle nodes: same seed -> same
digest across reruns and node types; different seed -> different digest.
- yambda_5b.gin: document precisely what $SEED controls and what it does
not (streaming data order, holdout split).
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: default INIT_CHECKSUM off (fp64 shard copy OOMs the build)
The startup reproducibility checksum computes per-shard sum/norm with
dtype=float64, which materializes a full fp64 copy of each local embedding
shard (>150 GiB for the big tables). After sharding leaves ~95 GiB of fp32
tables resident, that temporary leaves almost no HBM headroom and OOMs the
build during make_optimizer_and_shard on any node with residual memory.
Flip INIT_CHECKSUM to default 0 (opt-in) so normal launches and supervisor
resubmits never run it. Correct the utils.py + yambda_5b.gin comments that
claimed it had no full-size temporaries.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: add last_n UIH history strategy ($HISTORY_STRATEGY)
Adds a configurable Yambda UIH construction strategy alongside the existing
per-pool interleaved scheme. "last_n" takes the last HISTORY_LENGTH events of
any pool (listen+/like/skip) with no per-pool split, raising effective sequence
length (~2.7k -> ~4.1k) and letting the like share fall to its natural rate.
Default stays "interleaved" (no behavior change); strategy is resolved at
sample-build time so it reuses the existing on-disk cache (no rebuild).
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: logging-freeze prep — MIN_HISTORY=4086 default + HISTORY_STRATEGY passthrough
- gin: MIN_HISTORY default 0 -> 4086 (power-users floor at the full history
budget; maps to the existing positions_L4086.npy cache, no rebuild/no
shared-dir write). AUC_THRESHOLD left unchanged (0.80275) pending finalization.
- launch_slurm.sh: forward $HISTORY_STRATEGY through the worker docker exec -e
block (was silently dropped, so the knob never reached the worker); fix the
stale lr-override echo (gin default is 1e-7, not 0.001).
- README: document MIN_HISTORY default as 4086.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: untrack docs/v4_vs_v2_and_hstu_walkthrough.md
Stop tracking the local walkthrough doc (kept on disk, no longer in the repo).
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: scrub hardcoded username from reference comments
Replace the example username in launch_slurm.sh / streaming_resume_test.sh
comments with a generic <user> placeholder. Runtime defaults already derive
the container name + mounts from $USER, so only doc/example strings changed.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: README — add full single/multi-node reference run example
Document a complete sbatch launch (run-shape + data-fraction eval cadence) for
1-node and 2-node, noting the launchers differ only in --nodes.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: exclude eval/checkpoint overhead from step_ms timing
MetricsLogger now brackets eval and checkpoint phases with pause/resume
perf timers so the reported step_ms reflects pure train-step latency.
Adds wall_step_ms (inclusive), eval_ms, and ckpt_ms to the perf log line
and TensorBoard scalars (appended for parser backward-compat). Checkpoint
saves and eval windows are wrapped with categorized pause/resume calls.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: README — use AUC_THRESHOLD=0.80275 in example for gin-default consistency
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: make a bare sbatch reproduce the frozen reference run
orchestrate() now defaults to the reference run-shape (START_TS=0,
NUM_TRAIN_TS=299, full windows) + the data-fraction eval cadence
(EVAL_EVERY_DATA_PCT=0.005, per-window off), so `sbatch scripts/launch_slurm.sh`
needs no env knobs. SMOKE=1 restores the previous fast functional defaults
(short window, capped batches, per-window eval). The two eval cadences are
auto-deconflicted (explicit EVAL_EVERY_N_WINDOWS>0 disables data-pct). gin
library defaults + the resume/local smoke paths are unchanged. README updated to
the bare single/multi-node commands.
Co-authored-by: Cursor <cursoragent@cursor.com>
* dlrmv4: address PR review — mlperf_lo…
* Merge checklists in PR template * Delete .github/PULL_REQUEST_TEMPLATE directory * Add PR template title * Split PR templates and auto-append checklists from labels Use a short default template plus new-benchmark and new-round files. A pull_request_target workflow appends the matching checklist to the PR body so it applies for UI, CLI, and fork PRs. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: ShriyaRishab <spalsamudram@nvidia.com>
|
MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅ |
ShriyaRishab
added
the
new-benchmark
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
3 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.
Summary
Checklists are appended automatically from PR labels (works for UI, CLI, and fork PRs):
new-benchmark→ New Benchmark checklistnew-round→ New Round checklistAdd the matching label, or pick a template from the GitHub template picker /
gh pr create --template <file>.