Hard time replicating on 8xH100

[shenberg]

Hi,

I made what I think is a small improvement to the record (commit link) by developing on 1xH100, rebased on PR #159 and shelled out on an 8xH100 SXM bare-metal machine (Ubuntu 22 nvidia-docker) on Prime-Intellect to start collecting data. All of my results were roughly 0.7 seconds slower than I expected (I have 4 runs, 2 with my change and two at #159 and at first I thought my change just didn't improve run-time on 8xH100, until I thought to revert the code and run the original record version 😅).
Any advice on how to successfully replicate your results as a baseline, so I can show improvement? I want external validation 🙈
(and also to ask questions such as why do we even need that sa_lambda[0] multiply at all instead of just scaling the weights at init time (unsatisfying answer is because it degrades the results. Maybe the interplay between sa_lambda[0] being optimized by AdamW and vs QKVO by Muon matters, maybe it's the computation graph difference, I really don't know))

These are the baseline logs:

0d29bb84-83d1-446c-8aaf-c4deeb2df2d3.txt - 135780ms

f06653aa-edab-4ff4-bbc0-71a722af0a2b.txt - 135687ms

These are the modified version logs:

51b93080-55c3-4ad7-b27a-9784b8b71143.txt - 135088ms

07d52694-1cbd-49c5-9244-7a7883500fa5.txt - 135145ms

[ClassicLarry]

The last couple records have been ran on torch nightly 0926, which is a about 2 seconds faster than 1120. Can install via:
pip install torch==2.10.0.dev20250926+cu126 --index-url https://download.pytorch.org/whl/nightly/cu126

Different hardware will vary the runtime by 1-2s. To post a new record you don't need to beat the time, just show a speedup relative to your baseline, which I can validate carries over.

I like the idea behind your improvement!

[shenberg]

Thanks for the feedback!
I'll do a bunch of runs later today and update.

[shenberg]

So I did a bunch of runs but noticed an anomaly that makes the data somewhat unusable.
First of all, as the baseline I accidentally ran record 44 instead of your code. Big oops.

Worse, though, is that I had a strange, stocastic-but-semi-predictable noise pattern when running my code. Example plot of step times:

There are two unusual steps, at step 726 and 1023, that take around 0.5 seconds. I thought initially maybe I have some strange interaction with the window size change in YaRN (step 723-724), but this doesn't explain step 1023 and the lack of an issue around 1550. I ran around 15 runs and it's not always at the same spots, and it makes variance very high.

(The noise every 250 steps is due to the cuda synchronization around validation - step right after takes a hit, steps after it gain time due to the measurement simply being queuing work for cuda).

I did the same thing for the runs with record 44, and they didn't have this. I suspected it was due to data-loading being too slow maybe, and I was getting a warning about OMP_MAX_THREADS being set to 1 by torchrun so I thought maybe that's an issue, did a couple of runs with it set to 8 (machine had 104 cores so 8*8 seemed safe). This caused more noisy steps on my code and added noise to record 44 code. In short, AAAAAAAaaaahhhh.

Here's the same plot for one of your record runs (and I intentionally chose the noisiest one):

The plots for my replication of record 44 are similar. I'll try again next time I have a couple of hours.

[ClassicLarry]

That noise pattern/stalling has been observed before, but the root cause has not been identified. It depends on the specific hardware, the amount of computation happening, and the torch version. For instance if you pick different hardware it might start occurring on step 2152. As you observed, 2-3 specific spots in the training run will spike by 400ms. Last I checked it has nothing to do with when data is loaded, or the validation step, or the window size changes. It seems to be most sensitive to changes in the optimizers and their grad syncs. At some point I'll try and debug that with the profiler. #154 may have introduced this stalling back on some hardware due to the increase in Muon Optimizer time for Cautious Weight Decay.

My guess is, that if you observe stalling on your hardware for your submission, you will observe the same amount of stalling for the live record, so it won't interfere with showing that the new PR is faster. You'll probably need about 10 runs to show loss is under 3.28 with p<0.01. 3 runs of baseline is generally sufficient. I can help test if you run into issues.

[shenberg]

As a side-note, regarding your readme, I also ran a dense-connectivity experiment, using the skip connection mechanism (I'll get back to it) but I also made an attempt to add x0-residual mixing before the MLP (so lambda is 4 tensors per row). It improved convergence at a cost to runtime (Didn't try to see if I can reduce enough steps to make it worthwhile for the record).

Here are the weights of these scalars from a trained model (pre-attention x, pre-attention x0, pre-mlp x, pre-mlp x0):

tensor([[ 1.0000,  0.0000,  1.0000,  0.0000],
        [ 5.1552,  4.1552,  7.3093,  3.7690],
        [ 0.5204,  7.3022,  1.5996,  7.9686],
        [ 0.4849,  9.2451,  1.5034,  8.3961],
        [ 0.3342,  7.3814,  1.4735,  9.9009],
        [ 0.3087,  7.2016,  1.3673, 12.1176],
        [ 0.2748,  9.9410,  1.7711, 14.7112],
        [ 1.0000,  0.0000,  0.7980, 10.3982],
        [ 0.4198,  6.0632,  1.4630, 11.6521],
        [ 0.9300, 10.1884,  1.2211, 10.0390],
        [ 0.8024,  8.3456,  1.0307,  8.8656],
        [ 0.6096,  7.0954,  1.6096, 11.9012]])

The way I interpret the numbers is that attention influence is increased (residual scaled down before attention means attention output will be larger relative to the residual), and MLP influence is decreased. Behavior after layer 8 is different, I suspect due to it being chosen as the backoff layer (causality reversed relative to your hypothesis, I guess moving backoff to 9 and seeing the weights could each something about this).

I noticed the same special behavior for layer 8 when I did my dense-skip experiment. Here's a visualization of the connectivity matrix (row is the layer being calculated, column + 1 is the input layer, e.g. row 2 = layer 2 has a connection to layer 1 = col 0):

Layer 8 is the only stand-out that really mixes its preceding layers (I also tried letting layer 8 mix all preceding layers, again improving convergence but I'm pretty sure that not by enough to justify the runtime cost). (There's also something misleading about this because activation magnitudes grow very quickly through the layers so I think it would make sense to scale the values by the average magnitudes of activations in each layer in order to get a good sense of actual relative influence)

[ClassicLarry]

It could be worth testing a lambda on the mlp and attn output projections. You could apply the same method as your commit here to make it computationally efficient. Specifically the mlp out projection. Since that might enable similar benefits to what the lambda here on the value projection is doing.

[shenberg]

My original implementation was on the attn output projection because it seemed more natural: I did v = v + sa_lambda[0] * ve and then pre-multiplied O by sa_lambda[1] and set the sa_lambda init to (1.0, 0.5). Due to linearity in Vs it "should" behave the same. Anecdotally the final loss was worse, but maybe I was being misled by the randomness inherent in training (I only ran it once). Only once i saw the loss was worse did I move it to the QKV matrix. I rationalized the difference between the two as "the difference between the O and QKV versions is that the O version has a multiplicative interaction between sa_lambda[0] and [1], maybe that's more sensitive in the optimization process.

[shenberg]

Sigh. I got another machine, this time running as root on a spot-instance. I failed to replicate the result in two ways:

1. for some reason loss was worse so even 11 runs did not get significance (previous run attempt had like 0.001 significance so this is very strange)
2. I got better timing on your record(!) (3 runs, std, mean are (tensor(0.3286), tensor(134.0907)))

If I filtered by "good" runs on mine (=only one 400ms step) then I shave around 400ms from your run-time, but it seems as if my code triggers the issue pretty consistently. Seems like modifying the computation graph may be interacting with the distributed optimization.

[ClassicLarry]

To clarify, are you saying that after you made the change the p value was no longer under 0.01? Or that when you tried to replicate the old record it wasn’t below 0.01?

[shenberg]

My code with the original proposed change (sa_lambda[0] * qkv) I ran these
Test session 1: I ran 16 runs, got p=0.000 (loss mean 2.785, std 0.0011). Ran against the wrong version so don't have a relative timing comparison
Test session 2: 13 runs, p=0.099 (loss mean 3.2795, std 0.0013). Very noisy measurements, my times were 134.6415 +- 0.4375. Baseline (3 runs) times were 134.0907 +- 0.3286, loss 3.2788 +- 0.0018. My code having larger variance than the baseline with my code having 13 runs and baseline 3 is what frustrated me.

I decided to try some stuff with distributed Adam that I thought might, maybe, help but probably won't*. It didn't improve speed, but it showed me that everything was tied to the hardware I was using - I spun up a new instance and measurements were not as noisy.

This session had 15 non-baseline runs. I get good significance (p=0.0083, loss 3.279 +- 0.0014, time 135.3333 +- 0.3), however, 6 of those runs were me testing distributed Adam variants. If I discount those I get p=0.101, loss 3.2794 +- 0.0014, time 135.4751 +- 0.2779 (baseline (3 runs) time is 135.8643 +- 0.5710).

I find it hard to reconcile the extreme significance I got from my first session compared to later ones. Looking at sessions 2 & 3 I would have to conclude that I somehow hurt the final loss a little bit.

*I tried two things: 1. do the waits in step() according to the order of submission of reduce-scatter, 2. busy-wait on futures. I knew that it probably wouldn't matter because of everything being essentially barriers on the CUDA streams and not actually synchronous blocking

[shenberg]

Ok, I had some prime-intellect spot-instance time and I did some debugging:

I ran with NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=ALL and despite a mistake with logging (all ranks into one file, no timestamps), I managed to correlate these logs with the bad step:

1f89df06ce99:50400:50560 [6] NCCL INFO Received and initiated operation=Close res=011f89df06ce99:50400:50560 [6] NCCL INFO Received and initiated operation=Close res=011f89df06ce99:50400:50560 [6] NCCL INFO Received and initiated operation=Close res=0
1f89df06ce99:50400:50560 [6] NCCL INFO Received and initiated operation=Stop res=0
1f89df06ce99:50399:50553 [5] NCCL INFO [Proxy Service UDS] exit: stop 1 abortFlag 0
1f89df06ce99:50399:50399 [5] NCCL INFO NVLS Unbind MC handle 7fe5a3d7f7d0 size 2097152 dev 5
1f89df06ce99:50399:50399 [5] NCCL INFO NVLS Unmap mem UC handle 0x7fe5a3d7fd80(0x7fe5b7c00000) ucsize 2097152 MC handle 0x7fe5a3d7f7d0(0x7fe5b7e00000) mcsize 2097152
1f89df06ce99:50399:50399 [5] NCCL INFO comm 0x56333fbb8930 rank 5 nranks 8 cudaDev 5 busId bb000 - Destroy COMPLETE
1f89df01f89df06ce99:50396:50550 [0] NCCL INFO RAS current socket connection with 172.17.0.3<47475> closed by peer on receive; terminat1f89df01f89df06ce99:50396:50550 [0] NCCL INFO RAS handling local termination request
1f89df06ce99:50396:50550 [0] NCCL INFO RAS termin1f89df06ce99:50399:50545 [0] NCCL INFO RAS handling local termination request
1f89df06ce99:50399:50545 [0] NCCL INFO RAS terminating a connection with 172.17.0.3<60893>
1f89df06ce99:50399:50545 [0] NCCL INFO RAS terminating a connection with 172.17.0.3<52531>
1f89df06ce99:50399:50545 [0] NCCL INFO RAS thread terminating
ing
h 172.17.0.3<60893> closed by peer on receive; terminating it
1f89df06ce99:50401:50544 [0] NCCL INFO RAS current socket connection with 172.17.0.3<39245> closed by peer on receive; terminating it
1f89df06ce99:50401:50544 [0] NCCL INFO RAS handling local termination request
1f89df06ce99:50401:50544 [0] NCCL INFO RAS terminating a connection with 172.17.0.3<39245>
1f89df06ce99:50401:50544 [0] NCCL INFO RAS terminating a connection with 172.17.0.3<60893>
1f89df06ce99:50401:50544 [0] NCCL INFO RAS thread terminating
7f81d3c00000) ucsize 2097152 MC handle 0x7f81bfd7f7d0(0x7f81d3e00000) mcsize 2097152
1f89df06ce99:50395:50395 [1] NCCL INFO comm 0x55f252944d40 rank 1 nranks 8 cudaDev 1 busId 3b000 - Destroy COMPLETE
1f89df06ce99:50395:50547 [0] NCCL INFO RAS handling local termination request
1f89df06ce99:50395:50547 [0] NCCL INFO RAS terminating a connection with 172.17.0.3<49479>
1f89df06ce99:50395:50547 [0] NCCL INFO RAS terminating a connection with 172.17.0.3<39245>
1f89df06ce99:50395:50547 [0] NCCL INFO RAS thread terminating
..23}

This unusual, the rest of the file is a solid block of AllGather and ReduceScatter log lines (except for the init which I can also attach). What's stranger is that I would expect this to cause issues yet after these lines we return to AllGather and ReduceScatters as if nothing happened.

Anyhow, I was going to continue debugging this issue, and then I just got lucky: usually it would show up on the 2nd run, and now it just didn't so I continued to 11 runs (naturally the poor loss disappeared as well), and then ran the baseline, for what I believe is a record-setting run:

Note:
The run has another improvement (I think). I was staring at the profiling trace for the optimizer in perfetto:

and I noticed a gap between Adam and NorMuon, caused by waiting for DistAdam comms to finish before NorMuon starts, however they're independent - we only need the results from DistAdam before the start of the next step, and there's some time with no comms inside NorMuon, so I changed DistAdam to return the final future and only waited on it after the optimizer step(). I think this gave around 100ms of improvement out of the roughly 500ms. I decided to try testing only the lambda*w_qvko change before writing this just to tease out the separate contributions of the two changes, but I got a bad machine again. I decided I've donated enough money to Prime Intellect debugging this and left it at that.

The branch on my git is order_adam (sorry, I tried a bunch of things on that branch, I wasn't expecting this to be the run), forked from commit a842249.

Here are the logs:

logs.tar.gz

Extra note: I moved to using pytorch nightly 2.10.0.dev20251019+cu126 because the september version is not available anymore, I guess they keep versions only two months back. I also tried cu128 versions (some machines I used had CUDA 12.8 installed) and didn't make any difference in terms of runtime or of fixing poorly performing machine behavior

[ClassicLarry]

To have your change added you will need to open a PR against the most recent record (which was raised 11/29). Nightly 1204 is used in the final timing which seems pretty smooth. If you aren't planning to open a PR, someone else can probably roll your changes into their PR at some point.

[shenberg]

I've been trying to replicate again, and having issues. I found a tiny bugfix in the recent comms-compute overlap change for distadam, that I was sure also solved the noisy step problem (it doesn't). Pretty sure it shaves off a few hundred milliseconds by itself.