GMLP (blue) has slightly more params (200M vs. ~150M) than baseline (pink) and larger batch size - loss is plotted by tokens rather than steps to compensate. GMLP: 9x GMLP blocks + 3x local attention blocks Baseline: 12x dense attention blocks

pink = baseline blue = 512 90%, 2048 10% green = 128 50%, 2048 50%

Shortformer (https://arxiv.org/abs/2012.15832) style. 2 models trained on equal # tokens - one at seqlen = 512 for 90% of training, and seqlen = 2048 for the remaining 10%, and one at seqlen = 2048 for all of training. n.b - for some reason there's a bug in how wandb displays the runs, full screen and then minimize loss by tokens/wallclock time and it should display correctly. Also loss by tokens isn't 100% accurate since it scales by the initial seq length - but final losses should be equal # tokens

Shortformer (https://arxiv.org/abs/2012.15832) style. 2 models trained on equal number of tokens - one at seqlen = 512 for 90% of training, and seqlen = 2048 for the remaining 10%, and one at seqlen = 2048 for the whole training run. (the loss by tokens logging is somewhat inaccurate because it weights everything according to the initial seq length - but final number of tokens should be equal & hence final losses are comparable)

Small (150M) param model

on A100 x4 box

green = global blue = local (sliding window) pink = local (stepped)

dataset = enron emails

Results for rotary embeddings applied to only part of q/k. dim per head = 64 Pink - Learned Abs Baseline Brown - Rotary applied to 25% (16/64) Green - Rotary applied to 50% (32/64) Blue - Rotary applied to 100% (64/64) Other Pink - Rotary applied to 25% (16/64) every other layer

Results for rotary embeddings applied to only part of q/k. dim per head = 64 Pink - Learned Abs Baseline Brown - Rotary applied to 25% (16/64) Green - Rotary applied to 50% (32/64) Blue - Rotary applied to 100% (64/64)

Comparison between Ben's simplified RoPE implementation and the original one

150M param model on OWT2 with learned embeddings (blue) vs. rotary embeddings (green) vs. rpe (brown) vs. rpe with caching (peach)

150M param model on OWT2 with learned embeddings (blue) vs. rotary embeddings (green) vs. rpe (brown)

150M param model on OWT2 with learned embeddings (blue) vs. rotary embeddings (green)

(on A100s)

(model parallel only)

Comparison of hardware differences on Amazon P4DN instances, vs. an nvidia provided DGX box, vs. Coreweave machines.

Experiments using topk operator within self attention block

General scheme is mp = 2 (across nvlnk bridges), pp=n_gpus/2 (across all other connections), mb_size=16, g.a.s = 64. This may not be the ideal setup but should give us a general idea of the architecture's scalability. As we can see, the nvlink pairs topology appears to scale sub-linearly. It appears from the logs that the bottleneck seems to become the pipeline parallel connections as we scale.

16 GPUs (wandb says 8 but rank goes from 0 to 15), nhidden=1024, num_layers=24, sparsity on. Unusually low loss.

GPT2_XL_pipe, regular adam. mp=2, pp=2, size reduced to match 1-bit adam model.

GPT2-XL, regular adam, mp=2, pp=2, ZeRO-1

One-bit adam, no changes from the main branch

3D parallelism with no ZeRO, ran for 500 steps.