stitched together github repos into single long context input documents
nemotron v2 for web data
stack v3 for coding data (majority)
a lot of their synethetic data is from nvidia nemotron
taking rollouts from existing models to bootstrap agentic capabilities
for synthetic pretraining data they include both successful and failed rollouts and prepend the data with whether it failed or succeeded
the dataset is divided into categories (or buckets) like code, math, reasoning etc, and each bucket is given a data mixture weight
they have ~200 buckets (5 quality scores x 40 topics)
similar to regmix or olmix
at small scall, testing a data mix at 50bil tokens won’t accurately predict how the same data mix behaves when scaled to trillions of tokens (18trill)
so when running experiments on 50bil tokens they sub-sample
50billon/18trill = 1/360
Every dataset/bucket in the mix is downsampled by that exact ratio for the smaller experiment, instead of using all 50bil tokens out of the 18 tril from one/few bucket disproportionately
prevent multi-epoching (because that overfits)
Model Architecture
on B200 GQA and SWA is very fast, but linear attn was slower
GQA vs MLA
MLA + global attn on every layer but that increases quadratic attention cost
GQA ablations looked better
4 factors for picking arch
actual MFU on the hardware on the day you’d have to launch the hero run
what can you project the MFU to mid run
using EP kernels
as you do a scaling ladder for ~4 smaller scales, what is the loss across the scaling ladder and how to project that out to hero compute scale
checking for gradient spikes and loss spikes, and consider
they use hyperball
RMS norm after latent moe down proj, and RMS norm after the MoE block (which reduced loss by 30%)
Training
first GPU run, previously was training on TPUs with different arch decisions for TPU
initially trying to train with 600b params, and some amount of sparsity ratio by projecting isoFLOPs - turned out this size would not fit in a single gb200 rack
1 or 2 racks?
more sparsity is better - used 700B params and shard across 2 racks, with communication between fwd and bwd passes, but communication is slower
so 1 rack
how big of a model can you fit on 1 rack?
started with fsdp implementation (how much total model state can physically fit into the combined VRAM of a single rack without relying on slower, cross-rack communication.)
350B params
FSDP takes up memory because you need the matrix locally on the gpu to do expert fwd or bwd
EP allows to shard the experts across the gpus on the rack, and clears up more memory space
Expert parallelism - EP64 within a rack + FSDP
Using more expert granularity, with smaller experts but twice as more,
used along with EPA implementation, improved comm speed a bit more
trained with mixed precision - fp32 and bf16, wanted to be more cautious on this run so avoided fp8 or fp4
Kernels
EP
current implementation drops some tokens as it has fixed static shape
if more tokens are assigned to the expert than it handles, the tokens at the edge of the buffer gets dropped
Inference arch design
with a fixed amount of compute you can choose between spending that compute on more data or more params
whats the right param : data ratio to get best loss, and a fast model
13 times overtrained (trained on 13 times more [unique] data than chinchilla optimal)
inference efficiency
SWA on 3/4th of the layers
higher kv cache compression on the global layers (8:1 ration bw Q heads and KV heads)
MoE for more sparsity - gives better model without much inference cost
Hardware
considerations to being robust to hardware failure
had 12 NVL72s (12 racks and 72 B200s per rack)
but often that does not mean you have 12 * 72 B200s to work with bc of hardware failure
used only 11 racks
total gpu count = 64 x 11
the extra spare rack can be used to do moe kernel stuff while the hero run is running
SLA - service lease agreements
the marin team’s SLA was that 64 out of the 72 gpus must be up at any given time
snowball MoE - 67B-A2B MoE trained on 10T tokens, trained on TPUs
trying not to heavily rely on distillation
RL for both non agentic and agentic contexts
OPD - some teacher models to be considered are their 535B experts from early cooldowns of the hero run
Improvements / Next steps
if more time…
hybrid linear attention variant running efficiently as a next step
its hard to use off the shelf kernels or libraries unless you’re on a very similar setup as those implementations, often kernel libraries don’t give the promised perf because of different setup, different hardware, different model context, stack etc
