Marin 8B (aka Tootsie Roll run)

[dlwh]

Description

Tootsie Rolls famously use a graining process where the previous day's batch is folded into the next day's, resulting in consistency, or something.

We anticipate we'll be constantly improving our data mixture. We may as well get a main run started and improve the mixture as it progresses, rather than starting from scratch. By using WSD-S, we're basically just warm-starting anyway whenever we switch things up.

Hypothesis or Goal

1. Verify WSD-S works at this scale
2. Investigate folding in new data as this run progresses, using our best guess at what good data is as the run progresses.
3. (Maybe) is it better to fold in data: right after minima, or middle of cycle? My guess is mid-cycle.

Links

(Delete any that aren't applicable)

• WandB Report: https://wandb.ai/marin-community/marin/reports/Tootsie-8B---VmlldzoxMTY3MzU3OA
• Experiment JSON: https://marin.community/data-browser/experiment/?path=gs%3A//marin-us-central2/experiments/exp600_tootsie-4c95ae.json

Results

See comment stream for details.

[dlwh]

I have not been the best at updating this. In particular I made many fewer changes than initially planned.

I will attempt to go back and log what has happened:

840B: Increased Spacing for WSD-S:

After 840B Tokens, I switched from a WSD-S with cycle length 10k to cycle length 20k. (Decay period was still 10% of the cycle, so 1K steps became 2K steps). This improved performance on evals that look the most like the training data.

2.7T

!!! This has been superseded by a similar change documented in the next comment. We discovered that our copy of DCLM in us-central2 was corrupted whcih invalidated the run. In particular, we were actualyl already epoching, which is the cause of the sudden drop in loss.

At 2.7T Tokens, we made a lot of changes:

• We moved from 2 x v5e-256 to v4-2048
• 4M->12M token BS
• LR: 1e-3 -> 1.7e-3 (Following https://arxiv.org/abs/2205.10287 and increased batch size)
• Schedule: WSD-S -> Constant w/ 20% decay, EMA (Following https://arxiv.org/pdf/2412.19437 for EMA), 20% decay from a mix of vibes and experiments (TODO)
• Rope: Llama2 -> LLama 3 (oopsie)

Time will tell if this was a good idea

[dlwh]

2.7T Redux

We made the same changes again, except keeping it in europe for a while

• 4M->12M token BS
• LR: 1e-3 -> 1.7e-3 (Following https://arxiv.org/abs/2205.10287 and increased batch size)
• Schedule: WSD-S -> Constant w/ 20% decay, EMA (Following https://arxiv.org/pdf/2412.19437 for EMA), 20% decay from a mix of vibes and experiments (TODO)
• Rope: Llama2 -> LLama 3 (oopsie)

Now that we're not epoching the image is much more reasonable

[dlwh]

Cooldown Transition Take 1

We were coming up on the epoch boundary for DCLM so it was time to transition to mid-training/phase 2.

We attempted to do a dolmino-like cooldown based on some intuitions and guesses from #820, which was/is still ongoing. We decided to make some changes from Olmo 2's procedure, based on the differences in the pre-training mix and our intuitions.
In particular, we had an aversion to flan, due to it being too task-y and too synthetic. (One of us had prior bad experiences)

We elected to include:

Dataset	Value
dolma/pes2o	10.78
dolmino_dclm	67.77
finemath_3_plus	6.31
dolma/arxiv	5.17
dolma/stackexchange	3.17
starcoderdata	2.23
dolma/algebraic-stack	2.13
dolma/megawika	0.82
dolma/open-web-math	0.94
dolma/wiki	0.68

This was chosen based on 70% (DCLM-HQ + 0.1 starcoder) plus 30% HQ data (everything else) sampled proportional to token count.

We set a target of about 1.5 passes through the HQ data, which led us to 75k steps starting from step 740,500. We set a linear LR decay from the same point, decaying to .1 * max_lr (so 0.1 * 1.7e-3)

Our initial run is showing steep improvements in the scientific/HQ subsets of paloma, and a huge initial spike in c4en. The spike in c4en fell back down reasonably quickly.

Evals initially took a big hit. We'll see how it develops. (EDIT: Worked out!)

[dlwh]

Cooldown v1 Results

The result of this cooldown produced solid eval scores:

mmlu_0shot: 0.6247685514883919
mmlu_5shot: 0.6531833072211936
BBH few shot COT: 0.5702657041929043
GSM8K (8-shot): 0.5094768764215315
HumanEval (pass@1): 0.24390243902439024
MATH (4-shot): 0.1846
IFEval (prompt-level loose): 0.09242144177449169

MMLU 65.3 is better than DCLM and Olmo 2 7B

We're pretty far behind on math, but we didn't include any math in the previous cooldown.

Training curve:

Gray is the initial run which we did on the v5e-256s with WSD. the 3072 handoff (bs=12M tokens) worked reasonably well (phase 2, peach), and the cooldown (phase 3 green) saw a steep decline in trainin loss, reflecting both the train loss and the easier data.

Eval losses took a hit.

The EMA (dashed lines) in phase 2 worked quite well and tracked our WSD-S cooldowns well.

We saw a substantial jump in c4en during the cooldown, but this is expected given the sudden shift to a 70/30 mix.

[dlwh]

Dessert v1

While our numbers were pretty solid, we thought we could do better. Our ablations (#784) that we were running parallel suggested that FLAN was actually pretty important for task performance. In addition, as seen above our math scores weren't amazing. Olmo had a similar problem and solved it with their dolmino math sets.

Our model was already cooled down. What do? We decided to run what I called dessert (flan, get it?) where we just ran the model at the already cooled down rate (1.7e-4), which is reasonably high for a "cooled down" model. We'd run this for just a short little while.

The mix we went with was intended to be similar to the cooldown, but with the new datasets in: 0.7 (DCLM + 0.1 Starcoder) + .2 (Cooldown HQ) + .1 (Flan + Math)

Unfortunately, I ran 0.7 (DCLM + 0.1 Starcoder) + .1 (Cooldown HQ) + 0.2 (Flan + Math). Oops. Also some of the math didn't make it in.

Anyway, performance on MMLU wasn't better, and gsm8k even went down.

mmlu_5shot: 0.6486255519156815
mmlu_0shot: 0.6250534111949865
gsm8k (8-shot): 0.4799090219863533

To give you an idea of the LR schedule:

[dlwh]

Dessert v2

Ok, once more with feeling.

This time, we fixed math and the weights and started the run.

Not too much to report learning dynamics-wise, so let's look at the results:

mmlu 5 shot: 0.6511892892750321
mmlu 0 shot: 0.6249821962683378
gsm8k (8-shot): 0.6110689916603488
MATH: 0.2112

So GSM8k shot up (not surprising since the task is in the data) and MATH went up a bit too. MMLU was down slightly :(